Update app.py

app.py

@@ -1,10 +1,25 @@
-import os, io, json
+import os, json
 from typing import List, Dict, Any, Optional
 from PIL import Image
 import torch
 import gradio as gr
-import spaces
-from huggingface_hub import snapshot_download, HfHubHTTPError
+import spaces  # ZeroGPU decorator
+from huggingface_hub import snapshot_download
+
+# ----------------- Config (set in Space Secrets if private) -----------------
+# Your private repo that contains the base .safetensors and loras.json
+MODEL_REPO_ID = os.getenv("MODEL_REPO_ID", "DB2169/CyberPony_Lora").strip()
+# Exact filename of the base checkpoint inside the repo (case-sensitive)
+CHECKPOINT_FILENAME = os.getenv("CHECKPOINT_FILENAME", "SAFETENSORS_FILENAME.safetensors").strip()
+# Personal access token with read scope (required for private repos)
+HF_TOKEN = os.getenv("HF_TOKEN", None)
+# Toggle first-boot warmup (GPU-allocating on ZeroGPU)
+DO_WARMUP = os.getenv("WARMUP", "1") == "1"
+
+# Where snapshot_download will cache the repo
+REPO_DIR = "/home/user/model"
+
+# Supported schedulers
 from diffusers import (
     StableDiffusionXLPipeline,
     StableDiffusionPipeline,
@@ -15,14 +30,6 @@ from diffusers import (
     LMSDiscreteScheduler,
     PNDMScheduler,
 )
-
-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"
-
-REPO_DIR = "/home/user/model"
-
 SCHEDULERS = {
     "default": None,
     "euler_a": EulerAncestralDiscreteScheduler,
@@ -33,16 +40,23 @@ SCHEDULERS = {
     "dpmpp_2m": DPMSolverMultistepScheduler,
 }
 
+# Globals populated at startup
 pipe = None
 IS_SDXL = True
 LORA_MANIFEST: Dict[str, Dict[str, str]] = {}
-INIT_ERROR: Optional[str] = None  # expose bootstrap error to UI
+INIT_ERROR: Optional[str] = None
 
+# ----------------- Bootstrap (download + load on CPU) -----------------
 def bootstrap_model():
+    """
+    Downloads MODEL_REPO_ID into REPO_DIR and loads the single-file checkpoint.
+    Keeps pipeline on CPU; ZeroGPU attaches GPU inside the @spaces.GPU function.
+    """
     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."
+        INIT_ERROR = "Missing MODEL_REPO_ID or CHECKPOINT_FILENAME."
         print(f"[ERROR] {INIT_ERROR}")
         return
 
@@ -53,12 +67,8 @@ def bootstrap_model():
             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}"
+        INIT_ERROR = f"Failed to download repo {MODEL_REPO_ID}: {e}"
         print(f"[ERROR] {INIT_ERROR}")
         return
 
@@ -69,6 +79,7 @@ def bootstrap_model():
         return
 
     try:
+        # Try SDXL first
         _pipe = StableDiffusionXLPipeline.from_single_file(
             ckpt_path, torch_dtype=torch.float16, use_safetensors=True, add_watermarker=False
         )
@@ -84,6 +95,7 @@ def bootstrap_model():
             print(f"[ERROR] {INIT_ERROR}")
             return
 
+    # Light memory/perf tweaks
     if hasattr(_pipe, "enable_attention_slicing"):
         _pipe.enable_attention_slicing("max")
     if hasattr(_pipe, "enable_vae_slicing"):
@@ -91,6 +103,7 @@ def bootstrap_model():
     if hasattr(_pipe, "set_progress_bar_config"):
         _pipe.set_progress_bar_config(disable=True)
 
+    # Load LoRA manifest if present
     man_path = os.path.join(local_dir, "loras.json")
     manifest = {}
     if os.path.exists(man_path):
@@ -100,7 +113,7 @@ def bootstrap_model():
         except Exception as e:
             print(f"[WARN] Failed to parse loras.json: {e}")
 
-    # publish
+    # Publish globals
     global pipe, IS_SDXL, LORA_MANIFEST
     pipe = _pipe
     IS_SDXL = sdxl
@@ -125,6 +138,7 @@ def apply_loras(selected: List[str], scale: float, repo_dir: str):
     except Exception as e:
         print(f"[WARN] set_adapters failed: {e}")
 
+# ----------------- Generation (GPU-attached under ZeroGPU) -----------------
 @spaces.GPU
 def txt2img(
     prompt: str,
@@ -144,15 +158,16 @@ def txt2img(
         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)
 
+    # Optional scheduler switch
     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, REPO_DIR)
     if fuse_lora and loras:
         try:
@@ -182,8 +197,9 @@ def warmup():
     except Exception as e:
         print(f"[WARN] Warmup failed: {e}")
 
+# ----------------- UI -----------------
 with gr.Blocks(title="SDXL Space (ZeroGPU, single-file, LoRA-ready)") as demo:
-    status = gr.Markdown("")  # show init status/errors
+    status = gr.Markdown("")  # shows init result or errors
 
     with gr.Row():
         prompt = gr.Textbox(label="Prompt", lines=3)
@@ -206,7 +222,7 @@ with gr.Blocks(title="SDXL Space (ZeroGPU, single-file, LoRA-ready)") as demo:
     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", interactive=False)  # locked until model loads
+    btn = gr.Button("Generate", variant="primary", interactive=False)
     gallery = gr.Gallery(columns=4, height=420)
 
     def _startup():
@@ -230,4 +246,5 @@ with gr.Blocks(title="SDXL Space (ZeroGPU, single-file, LoRA-ready)") as demo:
         concurrency_id="gpu_queue",
     )
 
+# Gradio 4.x queue config (no deprecated args)
 demo.queue(max_size=32, default_concurrency_limit=1).launch()