Update app.py

app.py

@@ -1,76 +1,89 @@
-# app.py
+# app.py — SDXL + LoRA from Hugging Face (PEFT backend)
 import torch, gradio as gr
 from diffusers import StableDiffusionXLPipeline, DPMSolverMultistepScheduler
 
-# 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
+# --- Config ---
+BASE_MODEL = "stabilityai/stable-diffusion-xl-base-1.0"
+LORA_REPO  = "freshcodestech/LingoSpace"       # your HF LoRA repo
+LORA_FILE  = "LingoSpace-10.safetensors"       # exact filename in that repo
+ADAPTER    = "lingospace"                      # local nickname; can be any safe string
+LORA_W     = 0.8                               # LoRA strength (0.6–1.0 is typical)
 
-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 / dtype (GPU preferred)
 device = "cuda" if torch.cuda.is_available() else "cpu"
-dtype = torch.float16 if device == "cuda" else torch.float32
+dtype  = torch.float16 if device == "cuda" else torch.float32
 
-# Load SDXL (disable watermarker to avoid extra deps on Spaces)
+# --- Load SDXL base (disable watermarker to avoid extra deps on Spaces) ---
 pipe = StableDiffusionXLPipeline.from_pretrained(
-    BASE,
+    BASE_MODEL,
     torch_dtype=dtype,
     use_safetensors=True,
     add_watermarker=False,
-    **({"variant": "fp16"} if device == "cuda" else {})
+    **({"variant": "fp16"} if device == "cuda" else {})  # fp16 only on GPU
 )
 pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
 pipe = pipe.to(device)
 
-# Load LoRA from Hub
-pipe.load_lora_weights(LORA_REPO, weight_name=LORA_FILE)
+# --- Load LoRA from the Hub via PEFT backend ---
+# NOTE: This requires `peft` to be installed (see requirements.txt below).
+pipe.load_lora_weights(
+    LORA_REPO,
+    weight_name=LORA_FILE,
+    adapter_name=ADAPTER,
+    use_peft_backend=True
+)
 
-# If your diffusers supports adapter weights:
+# Preferred scaling API (if available)
 if hasattr(pipe, "set_adapters"):
-    # (adapter name is "default" when not specified in load_lora_weights)
-    pipe.set_adapters("default", weight=LORA_STRENGTH)
+    pipe.set_adapters(ADAPTER, weight=LORA_W)
 
-# Memory helpers
+# Memory helpers (helpful on T4 small)
 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()
 
-# 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 ---
+def infer(prompt, steps, cfg, h, w, lora_strength):
+    # Fallback strength control for older diffusers if set_adapters doesn’t exist:
+    kwargs = {}
+    if not hasattr(pipe, "set_adapters"):
+        kwargs["cross_attention_kwargs"] = {"scale": float(lora_strength)}
+    else:
+        # If adapter API exists, update weight dynamically
+        pipe.set_adapters(ADAPTER, weight=float(lora_strength))
+
+    out = pipe(
+        prompt,
+        num_inference_steps=int(steps),
+        guidance_scale=float(cfg),
+        height=int(h),
+        width=int(w),
+        **kwargs
+    )
+    return out.images[0]
+
+DEFAULT_PROMPT = "lingospace, cartoon, 2d illustration, flat colors, clean lineart, a raccoon in a compact utility space suit, holding a steaming bowl of rice, stars and soft nebula, cozy composition"
 
-def generate_images():
-    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
+# Use smaller defaults on CPU for speed
+DEF_STEPS = 30 if device == "cuda" else 15
+DEF_H = 1024 if device == "cuda" else 768
+DEF_W = 1024 if device == "cuda" else 768
 
 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: LingoSpace",
-    description=f"Device: {device.upper()} | Base: SDXL | LoRA: {LORA_REPO}/{LORA_FILE}",
+    fn=infer,
+    inputs=[
+        gr.Textbox(label="Prompt", value=DEFAULT_PROMPT, lines=3),
+        gr.Slider(10, 50, value=DEF_STEPS, step=1, label="Steps"),
+        gr.Slider(3.0, 9.0, value=7.5, step=0.5, label="CFG"),
+        gr.Slider(512, 1216, value=DEF_H, step=64, label="Height"),
+        gr.Slider(512, 1216, value=DEF_W, step=64, label="Width"),
+        gr.Slider(0.3, 1.2, value=LORA_W, step=0.05, label="LoRA Strength"),
+    ],
+    outputs=gr.Image(label="Output"),
+    title="SDXL + LoRA (freshcodestech/LingoSpace)",
+    description=f"Device: {device.upper()} | Base: {BASE_MODEL} | LoRA: {LORA_REPO}/{LORA_FILE}"
 )
 
 if __name__ == "__main__":