scscscc

app.py

@@ -1,18 +1,20 @@
 import os
 import torch
 import gradio as gr
-from diffusers import StableDiffusionImg2ImgPipeline
+from diffusers import DiffusionPipeline
 from PIL import Image
 
 # --- Configuration ---
-SPACE_TITLE = "🎨 Enhanced Studio Ghibli AI Art Generator"
-SPACE_DESCRIPTION = "Upload a portrait or a photo and transform it into a breathtaking Studio Ghibli-style masterpiece! Improved model and prompting for better results."
-MAIN_MODEL_ID = "nitrosocke/Ghibli-Diffusion"
+SPACE_TITLE = "🎨 Enhanced Studio Ghibli AI Art Generator (LoRA)"
+SPACE_DESCRIPTION = "Upload a portrait or a photo and transform it into a breathtaking Studio Ghibli-style masterpiece using a LoRA for fine-tuned results."
+BASE_MODEL_ID = "black-forest-labs/FLUX.1-dev"
+LORA_REPO_ID = "strangerzonehf/Flux-Ghibli-Art-LoRA"
+TRIGGER_WORD = "Ghibli Art"
 STRENGTH = 0.60  # Adjust for better balance between input and style
 GUIDANCE_SCALE = 7.5  # Increased for better prompt adherence
 NUM_INFERENCE_STEPS = 30  # Increased for potentially higher quality
 INPUT_IMAGE_SIZE = (512, 512)
-PROMPT_PREFIX = "Studio Ghibli anime-style illustration, "
+PROMPT_PREFIX = ""  # No need for separate prefix as LoRA is targeted
 NEGATIVE_PROMPT = "ugly, deformed, blurry, low quality, bad anatomy, bad proportions, disfigured, poorly drawn face, mutation, mutated, extra limbs, extra fingers, body horror, glitchy, tiling"
 
 # --- Device Setup ---
@@ -33,24 +35,26 @@ except Exception as e:
     print(f"⚠️ Torch initialization error: {e}")
 
 # --- Model Loading ---
-# Load the main Ghibli diffusion model for image-to-image
-img2img_pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
-    MAIN_MODEL_ID,
-    torch_dtype=torch.float16 if device == "cuda" else torch.float32,
-    use_safetensors=True,
-    low_cpu_mem_usage=True
-).to(device)
+pipe = DiffusionPipeline.from_pretrained(BASE_MODEL_ID, torch_dtype=torch.bfloat16)
+
+try:
+    pipe.load_lora_weights(LORA_REPO_ID)
+    print(f"✅ LoRA weights loaded from {LORA_REPO_ID}")
+except Exception as e:
+    print(f"⚠️ Error loading LoRA weights: {e}")
+
+pipe.to(device)
 
 # --- Optimization (Conditional for CUDA) ---
 if device == "cuda":
     try:
-        img2img_pipe.enable_xformers_memory_efficient_attention()
+        pipe.enable_xformers_memory_efficient_attention()
         print("✅ xFormers enabled!")
     except Exception as e:
         print(f"⚠️ xFormers not available: {e}")
-    img2img_pipe.enable_model_cpu_offload()
-img2img_pipe.enable_vae_slicing()  # Use img2img_pipe here
-img2img_pipe.enable_attention_slicing() # Use img2img_pipe here
+    pipe.enable_model_cpu_offload()
+pipe.enable_vae_slicing()
+pipe.enable_attention_slicing()
 
 # --- Image Transformation Function ---
 def transform_image(input_image):
@@ -60,16 +64,18 @@ def transform_image(input_image):
     try:
         input_image = input_image.resize(INPUT_IMAGE_SIZE)
 
-        output = img2img_pipe(
-            prompt=PROMPT_PREFIX + "portrait of a person", # Adjust prompt based on input
+        prompt = f"{PROMPT_PREFIX} {TRIGGER_WORD}, portrait of a person" # Incorporate trigger word
+
+        output = pipe(
+            prompt=prompt,
             image=input_image,
             strength=STRENGTH,
             guidance_scale=GUIDANCE_SCALE,
             num_inference_steps=NUM_INFERENCE_STEPS,
             negative_prompt=NEGATIVE_PROMPT,
-        )
+        ).images[0]
 
-        return output.images[0]
+        return output
     except Exception as e:
         print(f"❌ Error during image transformation: {e}")
         return None