Update app.py

app.py

@@ -1,67 +1,76 @@
-import torch
-import gradio as gr
+# app.py
+import torch, gradio as gr
 from diffusers import StableDiffusionXLPipeline, DPMSolverMultistepScheduler
 
-# Step 1: Define models
-base_model_sdxl = "stabilityai/stable-diffusion-xl-base-1.0"  
-lora_repo = "freshcodestech/LingoSpace"      
-lora_file = "LingoSpace-10.safetensors"      
+# Models
+BASE = "stabilityai/stable-diffusion-xl-base-1.0"
+LORA_REPO = "freshcodestech/LingoSpace"
+LORA_FILE = "LingoSpace-10.safetensors"
+LORA_STRENGTH = 0.8  # 0.6–1.0
 
-# Step 2: Load SDXL pipeline
+PROMPTS = [
+    "A cartoon raccoon in a space suit with some rice food, highly detailed, vibrant, futuristic art style",
+    "A cartoon bear in a space suit with some meat food, cinematic lighting, sci-fi background",
+    "A cartoon giraffe in a space suit with some pizza food, whimsical, colorful, space kitchen scene",
+]
+
+# Device / dtype
+device = "cuda" if torch.cuda.is_available() else "cpu"
+dtype = torch.float16 if device == "cuda" else torch.float32
+
+# Load SDXL (disable watermarker to avoid extra deps on Spaces)
 pipe = StableDiffusionXLPipeline.from_pretrained(
-    base_model_sdxl,
-    torch_dtype=torch.float16,
-    use_safetensors=True
+    BASE,
+    torch_dtype=dtype,
+    use_safetensors=True,
+    add_watermarker=False,
+    **({"variant": "fp16"} if device == "cuda" else {})
 )
-
-# Replace scheduler for better quality
 pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
+pipe = pipe.to(device)
 
-# Move to GPU
-pipe = pipe.to("cuda")
-
-# Step 3: Load LoRA weights
-pipe.load_lora_weights(lora_repo, weight_name=lora_file)
+# Load LoRA from Hub
+pipe.load_lora_weights(LORA_REPO, weight_name=LORA_FILE)
 
-# LoRA scaling factor (strength)
-lora_strength = 0.8   # Try 0.6–1.0
+# If your diffusers supports adapter weights:
+if hasattr(pipe, "set_adapters"):
+    # (adapter name is "default" when not specified in load_lora_weights)
+    pipe.set_adapters("default", weight=LORA_STRENGTH)
 
-# Step 4: Fixed prompts
-prompts = [
-    "A cartoon raccoon in a space suit with some rice food, highly detailed, vibrant, futuristic art style",
-    "A cartoon bear in a space suit with some meat food, cinematic lighting, sci-fi background",
-    "A cartoon giraffe in a space suit with some pizza food, whimsical, colorful, space kitchen scene"
-]
+# Memory helpers
+if device == "cuda":
+    if hasattr(pipe, "enable_xformers_memory_efficient_attention"):
+        pipe.enable_xformers_memory_efficient_attention()
+    pipe.enable_attention_slicing()
+    pipe.enable_vae_tiling()
 
-# Step 5: Generation settings
-num_steps = 30       # inference steps
-cfg = 7.5            # classifier-free guidance scale
-seed = 12345         # base seed for reproducibility
+# Gen settings (smaller & fewer steps on CPU so it doesn’t crawl)
+STEPS = 30 if device == "cuda" else 15
+CFG = 7.5
+H = 1024 if device == "cuda" else 768
+W = 1024 if device == "cuda" else 768
 
-# Inference function (returns list of images)
 def generate_images():
-    results = []
-    for idx, prompt in enumerate(prompts):
-        generator = torch.manual_seed(seed + idx)  # unique seed for each prompt
-        image = pipe(
-            prompt,
-            num_inference_steps=num_steps,
-            guidance_scale=cfg,
-            cross_attention_kwargs={"scale": lora_strength},
-            generator=generator,
-            height=1024,
-            width=1024
-        ).images[0]
-        results.append(image)
-    return results
+    images = []
+    for i, p in enumerate(PROMPTS):
+        out = pipe(
+            p,
+            num_inference_steps=STEPS,
+            guidance_scale=CFG,
+            height=H,
+            width=W,
+            # Works even if set_adapters isn't available (older diffusers):
+            cross_attention_kwargs={"scale": LORA_STRENGTH},
+        )
+        images.append(out.images[0])
+    return images
 
-# Step 6: Gradio UI
 demo = gr.Interface(
     fn=generate_images,
     inputs=None,
-    outputs=[gr.Image(label=f"Image {i+1}") for i in range(len(prompts))],
-    title="SDXL + LoRA: Cartoon Animals in Space Suits",
-    description="Generates fixed set of images using LoRA fine-tuning."
+    outputs=[gr.Image(label=f"Image {i+1}") for i in range(len(PROMPTS))],
+    title="SDXL + LoRA: LingoSpace",
+    description=f"Device: {device.upper()} | Base: SDXL | LoRA: {LORA_REPO}/{LORA_FILE}",
 )
 
 if __name__ == "__main__":