Update app.py

app.py

@@ -3,8 +3,8 @@ from typing import List, Dict, Any, Optional
 from PIL import Image
 import torch
 import gradio as gr
-import spaces  # ZeroGPU: decorate GPU-bound functions
-from huggingface_hub import snapshot_download
+import spaces
+from huggingface_hub import snapshot_download, HfHubHTTPError
 from diffusers import (
     StableDiffusionXLPipeline,
     StableDiffusionPipeline,
@@ -16,15 +16,12 @@ from diffusers import (
     PNDMScheduler,
 )
 
-# -------- Configuration (set as Space Secrets if needed) --------
-MODEL_REPO_ID = os.getenv("MODEL_REPO_ID", "DB2169/CyberPony_Lora")  # your model repo id
-CHECKPOINT_FILENAME = os.getenv("CHECKPOINT_FILENAME", "SAFETENSORS_FILENAME.safetensors")  # exact .safetensors name
-HF_TOKEN = os.getenv("HF_TOKEN", None)  # only required for private repos
-DO_WARMUP = os.getenv("WARMUP", "1") == "1"  # set to "0" to disable warmup
+MODEL_REPO_ID = os.getenv("MODEL_REPO_ID", "").strip()
+CHECKPOINT_FILENAME = os.getenv("CHECKPOINT_FILENAME", "").strip()
+HF_TOKEN = os.getenv("HF_TOKEN", None)
+DO_WARMUP = os.getenv("WARMUP", "1") == "1"
 
-# -------- Runtime defaults --------
-REPO_DIR = "/home/user/model"  # local cache mount for snapshot_download
-# Defer CUDA detection to GPU-run function for ZeroGPU; do not move to CUDA at import time
+REPO_DIR = "/home/user/model"
 
 SCHEDULERS = {
     "default": None,
@@ -36,24 +33,40 @@ SCHEDULERS = {
     "dpmpp_2m": DPMSolverMultistepScheduler,
 }
 
-# Globals populated on startup
 pipe = None
 IS_SDXL = True
 LORA_MANIFEST: Dict[str, Dict[str, str]] = {}
+INIT_ERROR: Optional[str] = None  # expose bootstrap error to UI
 
-
-# -------- Model bootstrap (CPU) --------
 def bootstrap_model():
-    global pipe, IS_SDXL, LORA_MANIFEST
-    local_dir = snapshot_download(
-        repo_id=MODEL_REPO_ID,
-        token=HF_TOKEN,
-        local_dir=REPO_DIR,
-        ignore_patterns=["*.md"],
-    )
+    global pipe, IS_SDXL, LORA_MANIFEST, INIT_ERROR
+    INIT_ERROR = None
+    if not MODEL_REPO_ID or not CHECKPOINT_FILENAME:
+        INIT_ERROR = "Missing MODEL_REPO_ID or CHECKPOINT_FILENAME environment variables."
+        print(f"[ERROR] {INIT_ERROR}")
+        return
+
+    try:
+        local_dir = snapshot_download(
+            repo_id=MODEL_REPO_ID,
+            token=HF_TOKEN,
+            local_dir=REPO_DIR,
+            ignore_patterns=["*.md"],
+        )
+    except HfHubHTTPError as e:
+        INIT_ERROR = f"Failed to download repo {MODEL_REPO_ID}: {e}"
+        print(f"[ERROR] {INIT_ERROR}")
+        return
+    except Exception as e:
+        INIT_ERROR = f"Unexpected error while downloading repo: {e}"
+        print(f"[ERROR] {INIT_ERROR}")
+        return
+
     ckpt_path = os.path.join(local_dir, CHECKPOINT_FILENAME)
     if not os.path.exists(ckpt_path):
-        raise FileNotFoundError(f"Checkpoint not found: {ckpt_path}")
+        INIT_ERROR = f"Checkpoint not found at {ckpt_path}. Check CHECKPOINT_FILENAME."
+        print(f"[ERROR] {INIT_ERROR}")
+        return
 
     try:
         _pipe = StableDiffusionXLPipeline.from_single_file(
@@ -61,12 +74,16 @@ def bootstrap_model():
         )
         sdxl = True
     except Exception:
-        _pipe = StableDiffusionPipeline.from_single_file(
-            ckpt_path, torch_dtype=torch.float16, use_safetensors=True
-        )
-        sdxl = False
+        try:
+            _pipe = StableDiffusionPipeline.from_single_file(
+                ckpt_path, torch_dtype=torch.float16, use_safetensors=True
+            )
+            sdxl = False
+        except Exception as e:
+            INIT_ERROR = f"Failed to load pipeline from {CHECKPOINT_FILENAME}: {e}"
+            print(f"[ERROR] {INIT_ERROR}")
+            return
 
-    # Keep on CPU until GPU-decorated call (ZeroGPU attaches GPU on demand)
     if hasattr(_pipe, "enable_attention_slicing"):
         _pipe.enable_attention_slicing("max")
     if hasattr(_pipe, "enable_vae_slicing"):
@@ -83,11 +100,12 @@ def bootstrap_model():
         except Exception as e:
             print(f"[WARN] Failed to parse loras.json: {e}")
 
+    # publish
+    global pipe, IS_SDXL, LORA_MANIFEST
     pipe = _pipe
     IS_SDXL = sdxl
     LORA_MANIFEST = manifest
 
-
 def apply_loras(selected: List[str], scale: float, repo_dir: str):
     if not selected or scale <= 0:
         return
@@ -107,8 +125,7 @@ def apply_loras(selected: List[str], scale: float, repo_dir: str):
     except Exception as e:
         print(f"[WARN] set_adapters failed: {e}")
 
-
-@spaces.GPU  # ZeroGPU: allocate/attach GPU for this function call
+@spaces.GPU
 def txt2img(
     prompt: str,
     negative: str,
@@ -123,19 +140,19 @@ def txt2img(
     lora_scale: float,
     fuse_lora: bool,
 ):
-    # Resolve device inside GPU context
+    if pipe is None:
+        raise RuntimeError(f"Model not initialized. {INIT_ERROR or 'Check Space secrets and logs.'}")
+
     local_device = "cuda" if torch.cuda.is_available() else "cpu"
     local_dtype = torch.float16 if local_device == "cuda" else torch.float32
     pipe.to(local_device)
 
-    # 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}")
 
-    # LoRAs
     apply_loras(loras, lora_scale, REPO_DIR)
     if fuse_lora and loras:
         try:
@@ -143,7 +160,6 @@ def txt2img(
         except Exception as e:
             print(f"[WARN] fuse_lora failed: {e}")
 
-    # Determinism
     generator = torch.Generator(device=local_device).manual_seed(int(seed)) if seed not in (None, "") else None
 
     kwargs: Dict[str, Any] = dict(
@@ -160,17 +176,14 @@ def txt2img(
         out = pipe(**kwargs)
     return out.images
 
-
 def warmup():
     try:
         _ = txt2img("warmup", "", 512, 512, 4, 4.0, 1, 1234, "default", [], 0.0, False)
     except Exception as e:
         print(f"[WARN] Warmup failed: {e}")
 
-
-# --------------------------- Build UI ---------------------------
-with gr.Blocks(title="SDXL Space (ZeroGPU, single-file checkpoint, LoRA-ready)") as demo:
-    gr.Markdown("### SDXL text‑to‑image with single‑file checkpoint and optional LoRAs")
+with gr.Blocks(title="SDXL Space (ZeroGPU, single-file, LoRA-ready)") as demo:
+    status = gr.Markdown("")  # show init status/errors
 
     with gr.Row():
         prompt = gr.Textbox(label="Prompt", lines=3)
@@ -193,21 +206,21 @@ with gr.Blocks(title="SDXL Space (ZeroGPU, single-file checkpoint, LoRA-ready)")
     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")
+    btn = gr.Button("Generate", variant="primary", interactive=False)  # locked until model loads
     gallery = gr.Gallery(columns=4, height=420)
 
-    # Load model + manifest, then populate LoRA choices
     def _startup():
         bootstrap_model()
-        return gr.CheckboxGroup.update(choices=list(LORA_MANIFEST.keys()))
+        if INIT_ERROR:
+            return gr.Markdown.update(value=f"❌ Init failed: {INIT_ERROR}"), gr.CheckboxGroup.update(choices=[]), gr.Button.update(interactive=False)
+        msg = f"✅ Model loaded from {MODEL_REPO_ID} ({'SDXL' if IS_SDXL else 'SD'})"
+        return gr.Markdown.update(value=msg), gr.CheckboxGroup.update(choices=list(LORA_MANIFEST.keys())), gr.Button.update(interactive=True)
 
-    demo.load(_startup, outputs=[lora_names])
+    demo.load(_startup, outputs=[status, lora_names, btn])
 
-    # Optional warmup (costs a tiny GPU run on first boot); set WARMUP=0 to skip
     if DO_WARMUP:
         demo.load(lambda: warmup(), inputs=None, outputs=None)
 
-    # Event binding inside Blocks; one GPU job at a time for SDXL
     btn.click(
         txt2img,
         inputs=[prompt, negative, width, height, steps, guidance, images, seed, scheduler, lora_names, lora_scale, fuse],
@@ -217,5 +230,4 @@ with gr.Blocks(title="SDXL Space (ZeroGPU, single-file checkpoint, LoRA-ready)")
         concurrency_id="gpu_queue",
     )
 
-# Global queue limits for Gradio 4.x
 demo.queue(max_size=32, default_concurrency_limit=1).launch()