Update app.py

app.py

@@ -1,9 +1,9 @@
-import os, io, json, base64
+import os, io, json
 from typing import List, Dict, Any, Optional
 from PIL import Image
 import torch
 import gradio as gr
-from huggingface_hub import snapshot_download
+from huggingface_hub import snapshot_download  # pulls your repo at startup
 from diffusers import (
     StableDiffusionXLPipeline,
     StableDiffusionPipeline,
@@ -15,11 +15,12 @@ from diffusers import (
     PNDMScheduler,
 )
 
-# ---- Config via Space Secrets / Environment ----
-MODEL_REPO_ID = os.getenv("MODEL_REPO_ID", "DB2169/CyberPony_Lora")
-CHECKPOINT_FILENAME = os.getenv("CHECKPOINT_FILENAME", "SAFETENSORS_FILENAME.safetensors")
-HF_TOKEN = os.getenv("HF_TOKEN", None)
+# -------- Configuration (set these in Space Secrets for private repos) --------
+MODEL_REPO_ID = os.getenv("MODEL_REPO_ID", "DB2169/CyberPony_Lora")          # e.g., your repo id
+CHECKPOINT_FILENAME = os.getenv("CHECKPOINT_FILENAME", "SAFETENSORS_FILENAME.safetensors")  # exact base ckpt filename
+HF_TOKEN = os.getenv("HF_TOKEN", None)                                        # optional if repo is public
 
+# -------- Runtime defaults --------
 device = "cuda" if torch.cuda.is_available() else "cpu"
 dtype = torch.float16 if device == "cuda" else torch.float32
 
@@ -33,64 +34,71 @@ SCHEDULERS = {
     "dpmpp_2m": DPMSolverMultistepScheduler,
 }
 
-# Globals populated at startup
+# Globals filled on startup
 pipe = None
 IS_SDXL = True
 LORA_MANIFEST: Dict[str, Dict[str, str]] = {}
+REPO_DIR = "/home/user/model"  # cached snapshot location in Spaces
 
+# -------- Model bootstrap --------
 def bootstrap_model():
     global pipe, IS_SDXL, LORA_MANIFEST
-
-    # Download your model repo (includes base .safetensors and loras.json)
-    repo_dir = snapshot_download(repo_id=MODEL_REPO_ID, token=HF_TOKEN, local_dir="/home/user/model", ignore_patterns=["*.md"])
-    ckpt_path = os.path.join(repo_dir, CHECKPOINT_FILENAME)
+    # Download/copy all repo files locally (weights + manifest)
+    local_dir = snapshot_download(
+        repo_id=MODEL_REPO_ID,
+        token=HF_TOKEN,
+        local_dir=REPO_DIR,
+        ignore_patterns=["*.md"],
+    )  # downloads your model repo into the container cache [web:362]
+
+    ckpt_path = os.path.join(local_dir, CHECKPOINT_FILENAME)
     if not os.path.exists(ckpt_path):
         raise FileNotFoundError(f"Checkpoint not found: {ckpt_path}")
 
-    # Load SDXL first; fallback to SD 1.x/2.x single-file
+    # Try SDXL single-file, then SD 1.x/2.x single-file
     try:
-        pipe_local = StableDiffusionXLPipeline.from_single_file(
+        _pipe = StableDiffusionXLPipeline.from_single_file(
             ckpt_path, torch_dtype=dtype, use_safetensors=True, add_watermarker=False
-        )
-        IS_SDXL = True
+        )  # SDXL loader [web:104]
+        sdxl = True
     except Exception:
-        pipe_local = StableDiffusionPipeline.from_single_file(
+        _pipe = StableDiffusionPipeline.from_single_file(
             ckpt_path, torch_dtype=dtype, use_safetensors=True
-        )
-        IS_SDXL = False
-
-    if hasattr(pipe_local, "enable_attention_slicing"):
-        pipe_local.enable_attention_slicing("max")
-    if hasattr(pipe_local, "enable_vae_slicing"):
-        pipe_local.enable_vae_slicing()
-    if hasattr(pipe_local, "set_progress_bar_config"):
-        pipe_local.set_progress_bar_config(disable=True)
-    pipe_local.to(device)
-
-    # Load LoRA manifest (optional)
-    man_path = os.path.join(repo_dir, "loras.json")
+        )  # SD 1.x/2.x fallback [web:104]
+        sdxl = False
+
+    if hasattr(_pipe, "enable_attention_slicing"):
+        _pipe.enable_attention_slicing("max")
+    if hasattr(_pipe, "enable_vae_slicing"):
+        _pipe.enable_vae_slicing()
+    if hasattr(_pipe, "set_progress_bar_config"):
+        _pipe.set_progress_bar_config(disable=True)
+    _pipe.to(device)
+
+    # Load LoRA manifest if present
+    man_path = os.path.join(local_dir, "loras.json")
+    manifest = {}
     if os.path.exists(man_path):
         try:
             with open(man_path, "r", encoding="utf-8") as f:
-                LORA_MANIFEST = json.load(f)
-        except Exception:
-            LORA_MANIFEST = {}
-    else:
-        LORA_MANIFEST = {}
+                manifest = json.load(f)
+        except Exception as e:
+            print(f"[WARN] Failed to parse loras.json: {e}")
 
-    return pipe_local
+    # Publish globals
+    return _pipe, sdxl, manifest
 
 def apply_loras(selected: List[str], scale: float):
     if not selected or scale <= 0:
         return
+    # Each selected LoRA should exist in manifest; supports repo/weight_name or local 'path'
     for name in selected:
         meta = LORA_MANIFEST.get(name)
         if not meta:
             continue
         try:
             if "path" in meta:
-                # If you later store LoRA files inside the model repo
-                pipe.load_lora_weights(os.path.join("/home/user/model", meta["path"]), adapter_name=name)
+                pipe.load_lora_weights(os.path.join(REPO_DIR, meta["path"]), adapter_name=name)
             else:
                 pipe.load_lora_weights(meta.get("repo", ""), weight_name=meta.get("weight_name"), adapter_name=name)
         except Exception as e:
@@ -114,12 +122,14 @@ def txt2img(
     lora_scale: float,
     fuse_lora: bool,
 ):
-    if scheduler and scheduler in SCHEDULERS and SCHEDULERS[scheduler] is not None:
+    # Scheduler swap
+    if scheduler in SCHEDULERS and SCHEDULERS[scheduler] is not None:
         try:
             pipe.scheduler = SCHEDULERS[scheduler].from_config(pipe.scheduler.config)
         except Exception as e:
             print(f"[WARN] Scheduler switch failed: {e}")
 
+    # Apply LoRAs
     apply_loras(loras, lora_scale)
     if fuse_lora and loras:
         try:
@@ -127,7 +137,9 @@ def txt2img(
         except Exception as e:
             print(f"[WARN] fuse_lora failed: {e}")
 
+    # Determinism
     generator = torch.Generator(device=device).manual_seed(int(seed)) if seed not in (None, "") else None
+
     kwargs: Dict[str, Any] = dict(
         prompt=prompt or "",
         negative_prompt=negative or None,
@@ -142,18 +154,15 @@ def txt2img(
     return out.images
 
 def warmup():
-    # Tiny, fast pass to initialize CUDA kernels and weight graphs
+    # Small, fast call to initialize kernels/graphs so first user is instant
     try:
-        _ = txt2img(
-            "warmup", "", 512 if IS_SDXL else 512, 512 if IS_SDXL else 512,
-            5, 5.0, 1, 1234, "default", [], 0.0, False
-        )
+        _ = txt2img("warmup", "", 512, 512, 4, 4.0, 1, 1234, "default", [], 0.0, False)
     except Exception as e:
-        print(f"[WARN] warmup failed: {e}")
+        print(f"[WARN] Warmup failed: {e}")
 
-# Build UI
-with gr.Blocks(title="SDXL Space (runs Diffusers directly)") as demo:
-    gr.Markdown("### SDXL text‑to‑image (single‑file checkpoint) with optional LoRAs")
+# --------------------------- Build the UI inside Blocks ---------------------------
+with gr.Blocks(title="SDXL Space (single-file, LoRA-ready)") as demo:  # Blocks context required for events [web:371]
+    gr.Markdown("### SDXL text‑to‑image (single‑file checkpoint) with optional LoRAs")  # UI heading [web:147]
     with gr.Row():
         prompt = gr.Textbox(label="Prompt", lines=3)
         negative = gr.Textbox(label="Negative Prompt", lines=3)
@@ -168,36 +177,33 @@ with gr.Blocks(title="SDXL Space (runs Diffusers directly)") as demo:
         seed = gr.Number(value=None, precision=0, label="Seed (blank=random)")
         scheduler = gr.Dropdown(list(SCHEDULERS.keys()), value="dpmpp_2m", label="Scheduler")
 
-    # LoRA controls: multi-select from loras.json
-    lora_names = gr.CheckboxGroup(choices=[], label="LoRAs (from loras.json; select any)")
+    # LoRA multi-select populated after manifest loads
+    lora_names = gr.CheckboxGroup(choices=[], label="LoRAs (from loras.json)")
     lora_scale = gr.Slider(0.0, 1.5, 0.7, step=0.05, label="LoRA scale")
     fuse = gr.Checkbox(label="Fuse LoRA (faster after load)")
 
     btn = gr.Button("Generate", variant="primary")
     gallery = gr.Gallery(columns=4, height=420)
 
+    # Startup loader (runs at app load)
     def _startup():
-        global pipe
-        pipe = bootstrap_model()
-        # Fill LoRA choices after manifest loads
+        global pipe, IS_SDXL, LORA_MANIFEST
+        pipe, IS_SDXL, LORA_MANIFEST = bootstrap_model()
         return gr.CheckboxGroup.update(choices=list(LORA_MANIFEST.keys()))
-    demo.load(_startup, outputs=[lora_names])
-    demo.load(lambda: warmup(), inputs=None, outputs=None)
+    demo.load(_startup, outputs=[lora_names])  # fill LoRA list once model is ready [web:147]
 
+    # Warm-up pass after model load for snappy first request
+    demo.load(lambda: warmup(), inputs=None, outputs=None)  # performance warmup [web:356]
+
+    # Wire the button click inside Blocks, with per-event concurrency control
     btn.click(
         txt2img,
         inputs=[prompt, negative, width, height, steps, guidance, images, seed, scheduler, lora_names, lora_scale, fuse],
         outputs=[gallery],
         api_name="txt2img",
-    )
+        concurrency_limit=1,          # one GPU job at a time for SDXL
+        concurrency_id="gpu_queue",   # shared queue id if you add more GPU events
+    )  # per-event queue parameters in Gradio 4.x [web:388][web:373]
 
-# Tune queue for Spaces (concurrency + backlog)
-btn.click(
-    txt2img,
-    inputs=[prompt, negative, width, height, steps, guidance, images, seed, scheduler, lora_names, lora_scale, fuse],
-    outputs=[gallery],
-    api_name="txt2img",
-    concurrency_limit=1,          # keep 1 GPU job at a time
-    concurrency_id="gpu_queue"    # share queue if you add more GPU events later
-)
-demo.queue(max_size=32, default_concurrency_limit=1).launch()
+# Global queue config (no deprecated args)
+demo.queue(max_size=32, default_concurrency_limit=1).launch()  # supported queue pattern in Gradio 4.x [web:373][web:381]