Update app.py

app.py

@@ -6,30 +6,52 @@ from PIL import Image
 # ✅ Base model (commercial use allowed)
 BASE_MODEL = "stabilityai/stable-diffusion-xl-base-1.0"
 
-# LoRAs
-LORA_1 = "gh1bli-style.safetensors"
-LORA_2 = "ghibli_landscape_lora.safetensors"
+# ✅ Local LoRA weights
+LORA_1 = "./gh1bli-style.safetensors"
+LORA_2 = "./ghibli_landscape_lora.safetensors"
 
+# ✅ Device setup
 device = "cuda" if torch.cuda.is_available() else "cpu"
 dtype = torch.float16 if torch.cuda.is_available() else torch.float32
 
-# 🧠 Load base text-to-image pipeline
+# 🧠 Load base pipeline
 print("🔹 Loading SDXL base model...")
 pipe_txt2img = DiffusionPipeline.from_pretrained(
     BASE_MODEL,
-    torch_dtype=dtype,
+    dtype=dtype,
     use_safetensors=True,
 ).to(device)
 
-# 🧩 Apply both LoRAs
+# ✅ Enable CPU/GPU memory optimization
+if device == "cuda":
+    print("🚀 Using GPU optimization")
+    pipe_txt2img.enable_model_cpu_offload()  # For big SDXL weights
+else:
+    print("🧩 Using CPU memory optimization")
+    pipe_txt2img.enable_attention_slicing()  # Reduce RAM spikes
+    pipe_txt2img.enable_sequential_cpu_offload()
+
+# 🧩 Load both LoRAs (PEFT-compatible)
 print("🔹 Applying Ghibli-style LoRAs...")
-pipe_txt2img.load_lora_weights(LORA_1)
-pipe_txt2img.load_lora_weights(LORA_2)
+pipe_txt2img.load_lora_weights(LORA_1, adapter_name="ghibli_style")
+pipe_txt2img.load_lora_weights(LORA_2, adapter_name="ghibli_landscape")
+
+# 🧠 Merge both styles
+pipe_txt2img.set_adapters(["ghibli_style", "ghibli_landscape"], adapter_weights=[0.7, 0.6])
 
-# 🖼️ Image-to-image pipeline (same model + LoRAs)
+# 🖼️ Image-to-Image pipeline (inherits adapters)
 print("🔹 Setting up image-to-image pipeline...")
 pipe_img2img = StableDiffusionXLImg2ImgPipeline(**pipe_txt2img.components)
+pipe_img2img.set_adapters(["ghibli_style", "ghibli_landscape"], adapter_weights=[0.7, 0.6])
+
+# ✅ Add same memory optimization for img2img
+if device == "cuda":
+    pipe_img2img.enable_model_cpu_offload()
+else:
+    pipe_img2img.enable_attention_slicing()
+    pipe_img2img.enable_sequential_cpu_offload()
 
+# 🎨 Generation function
 def generate(prompt, steps=30, guidance=7.5, seed=42, strength=0.6, image=None):
     generator = torch.Generator(device=device).manual_seed(int(seed))
     
@@ -55,7 +77,7 @@ def generate(prompt, steps=30, guidance=7.5, seed=42, strength=0.6, image=None):
     
     return result
 
-# 🎨 Gradio Interface
+# 🎛️ Gradio UI
 demo = gr.Interface(
     fn=generate,
     inputs=[
@@ -67,8 +89,8 @@ demo = gr.Interface(
         gr.Image(label="Upload Image (optional)", type="filepath"),
     ],
     outputs=gr.Image(label="Generated Image"),
-    title="Ghibli Style Maker – Text & Image to Image",
-    description="Create Ghibli-style art from a text prompt or transform any photo into Ghibli-inspired scenery using SDXL + LoRAs.",
+    title="🎨 Ghibli Style Maker – Text & Image to Image",
+    description="Generate or transform images in Studio Ghibli-inspired style using SDXL and LoRAs.",
 )
 
 if __name__ == "__main__":