Update app.py

app.py

@@ -1,114 +1,68 @@
 import gradio as gr
 import torch
-from diffusers import DiffusionPipeline, StableDiffusionImg2ImgPipeline
+from diffusers import StableDiffusionXLImg2ImgPipeline
 from safetensors.torch import load_file
 from PIL import Image
-import os
 
-# Your LoRA files
-LORA_FILES = [
-    "gh1bli-style.safetensors",
-    "ghibli_landscape_lora.safetensors",
-]
+# --- Base SDXL model ---
+BASE_MODEL = "stabilityai/stable-diffusion-xl-base-1.0"
+LORA_PATH = "studioghibli_flux_r32-v2.safetensors"
 
-# Function to check LoRA type
-def detect_lora_type(lora_path):
-    try:
-        keys = load_file(lora_path).keys()
-        # SDXL LoRAs have transformer_blocks, SD1.5 ones don’t
-        if any("transformer_blocks" in k for k in keys):
-            return "SDXL"
-        return "SD1.5"
-    except Exception as e:
-        print(f"⚠️ Could not read {lora_path}: {e}")
-        return "UNKNOWN"
-
-# Detect which model type to use
-detected_type = None
-for lora in LORA_FILES:
-    if os.path.exists(lora):
-        t = detect_lora_type(lora)
-        print(f"🔍 Detected {lora} → {t}")
-        if t != "UNKNOWN":
-            detected_type = t
-            break
-
-# Fallback if nothing detected
-if detected_type is None:
-    detected_type = "SD1.5"
-
-# Choose model accordingly
-if detected_type == "SDXL":
-    BASE_MODEL = "stabilityai/stable-diffusion-xl-base-1.0"
-    print("✅ Using SDXL base model")
-else:
-    BASE_MODEL = "runwayml/stable-diffusion-v1-5"
-    print("✅ Using SD1.5 base model")
-
-# Load base model
+# --- Setup device & dtype ---
 device = "cuda" if torch.cuda.is_available() else "cpu"
 dtype = torch.float16 if torch.cuda.is_available() else torch.float32
 
-pipe_txt2img = DiffusionPipeline.from_pretrained(
+print("🔹 Loading SDXL base model...")
+pipe = StableDiffusionXLImg2ImgPipeline.from_pretrained(
     BASE_MODEL,
     torch_dtype=dtype,
     use_safetensors=True,
+    variant="fp16" if torch.cuda.is_available() else None,
 ).to(device)
 
-# Apply LoRAs safely
-for lora in LORA_FILES:
-    if os.path.exists(lora):
-        try:
-            print(f"🎨 Loading LoRA: {lora}")
-            pipe_txt2img.load_lora_weights(lora)
-        except Exception as e:
-            print(f"⚠️ Skipped {lora}: {e}")
-
-# Image-to-Image (if supported)
-if detected_type == "SDXL":
-    from diffusers import StableDiffusionXLImg2ImgPipeline
-    pipe_img2img = StableDiffusionXLImg2ImgPipeline(**pipe_txt2img.components)
-else:
-    pipe_img2img = StableDiffusionImg2ImgPipeline(**pipe_txt2img.components)
+# --- Apply LoRA weights ---
+print("🎨 Applying Ghibli-style LoRA...")
+try:
+    lora_weights = load_file(LORA_PATH)
+    pipe.unet.load_state_dict(lora_weights, strict=False)
+    print("✅ LoRA loaded successfully.")
+except Exception as e:
+    print(f"⚠️ Failed to load LoRA: {e}")
 
-# Generate 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))
-    
-    if image is not None:
-        init_image = Image.open(image).convert("RGB").resize((768, 768))
-        result = pipe_img2img(
-            prompt=prompt,
-            image=init_image,
-            strength=float(strength),
-            num_inference_steps=int(steps),
-            guidance_scale=float(guidance),
-            generator=generator,
-        ).images[0]
-    else:
-        result = pipe_txt2img(
-            prompt=prompt,
-            num_inference_steps=int(steps),
-            guidance_scale=float(guidance),
-            generator=generator,
-        ).images[0]
+# --- Ghibli-style conversion ---
+def ghibli_style(image, steps=30, guidance=7.5, strength=0.6, seed=42):
+    if image is None:
+        raise gr.Error("Please upload an image to convert.")
     
+    generator = torch.Generator(device=device).manual_seed(int(seed))
+    init_image = Image.open(image).convert("RGB").resize((1024, 1024))
+
+    prompt = "Ghibli-style art, soft lighting, painterly textures, cinematic color palette"
+
+    result = pipe(
+        prompt=prompt,
+        image=init_image,
+        strength=float(strength),
+        num_inference_steps=int(steps),
+        guidance_scale=float(guidance),
+        generator=generator,
+    ).images[0]
+
     return result
 
-# Gradio UI
+# --- Gradio Interface ---
 demo = gr.Interface(
-    fn=generate,
+    fn=ghibli_style,
     inputs=[
-        gr.Textbox(label="Prompt", placeholder="A Ghibli-style mountain village at sunset"),
+        gr.Image(label="Upload Image", type="filepath"),
         gr.Slider(10, 50, 30, step=1, label="Inference Steps"),
         gr.Slider(1, 15, 7.5, step=0.5, label="Guidance Scale"),
+        gr.Slider(0.1, 1.0, 0.6, step=0.1, label="Style Strength"),
         gr.Number(label="Seed", value=42),
-        gr.Slider(0.1, 1.0, 0.6, step=0.1, label="Strength (for image input)"),
-        gr.Image(label="Upload Image (optional)", type="filepath"),
     ],
-    outputs=gr.Image(label="Generated Image"),
-    title="Ghibli Style Maker – Auto Model Switch",
-    description="Automatically detects if your LoRA is for SD1.5 or SDXL and generates Studio Ghibli–style art from text or image.",
+    outputs=gr.Image(label="Ghibli Style Output"),
+    title="Ghibli Style Image Converter",
+    description="Upload any image and transform it into a Studio Ghibli-style artwork using the Flux LoRA and SDXL model.",
 )
 
 if __name__ == "__main__":