Update app.py

app.py

@@ -3,23 +3,8 @@ from typing import List, Dict, Any, Optional
 from PIL import Image
 import torch
 import gradio as gr
-import spaces  # ZeroGPU decorator
+import spaces
 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,
@@ -30,6 +15,19 @@ from diffusers import (
     LMSDiscreteScheduler,
     PNDMScheduler,
 )
+
+# ----------------- Config (set in Space Secrets if private) -----------------
+MODEL_REPO_ID = os.getenv("MODEL_REPO_ID", "DB2169/CyberPony_Lora").strip()
+CHECKPOINT_FILENAME = os.getenv("CHECKPOINT_FILENAME", "SAFETENSORS_FILENAME.safetensors").strip()
+HF_TOKEN = os.getenv("HF_TOKEN", None)
+DO_WARMUP = os.getenv("WARMUP", "1") == "1"   # set WARMUP=0 to skip the first warmup call
+
+# Optional override: JSON string for LoRA manifest (same shape as loras.json)
+LORAS_JSON = os.getenv("LORAS_JSON", "").strip()
+
+# Where snapshot_download caches the repo in the container
+REPO_DIR = "/home/user/model"
+
 SCHEDULERS = {
     "default": None,
     "euler_a": EulerAncestralDiscreteScheduler,
@@ -46,11 +44,60 @@ IS_SDXL = True
 LORA_MANIFEST: Dict[str, Dict[str, str]] = {}
 INIT_ERROR: Optional[str] = None
 
+# ----------------- Helpers -----------------
+def load_lora_manifest(repo_dir: str) -> Dict[str, Dict[str, str]]:
+    """
+    Manifest load order:
+    1) Environment variable LORAS_JSON (if provided)
+    2) loras.json inside the downloaded model repo
+    3) loras.json at the Space root (next to app.py)
+    4) Built-in fallback with MoriiMee_Gothic you provided
+    """
+    # 1) From env JSON
+    if LORAS_JSON:
+        try:
+            parsed = json.loads(LORAS_JSON)
+            if isinstance(parsed, dict):
+                return parsed
+        except Exception as e:
+            print(f"[WARN] Failed to parse LORAS_JSON: {e}")
+
+    # 2) From repo
+    repo_manifest = os.path.join(repo_dir, "loras.json")
+    if os.path.exists(repo_manifest):
+        try:
+            with open(repo_manifest, "r", encoding="utf-8") as f:
+                parsed = json.load(f)
+            if isinstance(parsed, dict):
+                return parsed
+        except Exception as e:
+            print(f"[WARN] Failed to parse repo loras.json: {e}")
+
+    # 3) From Space root
+    local_manifest = os.path.join(os.getcwd(), "loras.json")
+    if os.path.exists(local_manifest):
+        try:
+            with open(local_manifest, "r", encoding="utf-8") as f:
+                parsed = json.load(f)
+            if isinstance(parsed, dict):
+                return parsed
+        except Exception as e:
+            print(f"[WARN] Failed to parse local loras.json: {e}")
+
+    # 4) Built-in fallback: your MoriiMee Gothic LoRA
+    print("[INFO] Using built-in LoRA fallback manifest.")
+    return {
+        "MoriiMee_Gothic": {
+            "repo": "LyliaEngine/MoriiMee_Gothic_Niji_Style_Illustrious_r1",
+            "weight_name": "MoriiMee_Gothic_Niji_Style_Illustrious_r1.safetensors"
+        }
+    }
+
 # ----------------- 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.
+    Downloads MODEL_REPO_ID into REPO_DIR and loads the single-file checkpoint,
+    keeping weights on CPU; ZeroGPU attaches GPU only inside @spaces.GPU calls.
     """
     global pipe, IS_SDXL, LORA_MANIFEST, INIT_ERROR
     INIT_ERROR = None
@@ -79,7 +126,7 @@ def bootstrap_model():
         return
 
     try:
-        # Try SDXL first
+        # Attempt SDXL first (text_encoder_2 present)
         _pipe = StableDiffusionXLPipeline.from_single_file(
             ckpt_path, torch_dtype=torch.float16, use_safetensors=True, add_watermarker=False
         )
@@ -91,11 +138,10 @@ def bootstrap_model():
             )
             sdxl = False
         except Exception as e:
-            INIT_ERROR = f"Failed to load pipeline from {CHECKPOINT_FILENAME}: {e}"
+            INIT_ERROR = f"Failed to load pipeline: {e}"
             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"):
@@ -103,18 +149,10 @@ 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):
-        try:
-            with open(man_path, "r", encoding="utf-8") as f:
-                manifest = json.load(f)
-        except Exception as e:
-            print(f"[WARN] Failed to parse loras.json: {e}")
+    manifest = load_lora_manifest(local_dir)
+    print(f"[INFO] LoRAs available: {list(manifest.keys())}")
 
-    # Publish globals
-    global pipe, IS_SDXL, LORA_MANIFEST
+    # Publish
     pipe = _pipe
     IS_SDXL = sdxl
     LORA_MANIFEST = manifest
@@ -125,20 +163,23 @@ def apply_loras(selected: List[str], scale: float, repo_dir: str):
     for name in selected:
         meta = LORA_MANIFEST.get(name)
         if not meta:
+            print(f"[WARN] Requested LoRA '{name}' not in manifest.")
             continue
         try:
             if "path" in meta:
                 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)
+            print(f"[INFO] Loaded LoRA: {name}")
         except Exception as e:
             print(f"[WARN] LoRA load failed for {name}: {e}")
     try:
         pipe.set_adapters(selected, adapter_weights=[float(scale)] * len(selected))
+        print(f"[INFO] Activated LoRAs: {selected} at scale {scale}")
     except Exception as e:
         print(f"[WARN] set_adapters failed: {e}")
 
-# ----------------- Generation (GPU-attached under ZeroGPU) -----------------
+# ----------------- Generation (ZeroGPU) -----------------
 @spaces.GPU
 def txt2img(
     prompt: str,
@@ -160,14 +201,12 @@ def txt2img(
     local_device = "cuda" if torch.cuda.is_available() else "cpu"
     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:
@@ -199,7 +238,7 @@ def warmup():
 
 # ----------------- UI -----------------
 with gr.Blocks(title="SDXL Space (ZeroGPU, single-file, LoRA-ready)") as demo:
-    status = gr.Markdown("")  # shows init result or errors
+    status = gr.Markdown("")
 
     with gr.Row():
         prompt = gr.Textbox(label="Prompt", lines=3)
@@ -228,9 +267,10 @@ with gr.Blocks(title="SDXL Space (ZeroGPU, single-file, LoRA-ready)") as demo:
     def _startup():
         bootstrap_model()
         if INIT_ERROR:
-            return gr.Markdown.update(value=f"❌ Init failed: {INIT_ERROR}"), gr.CheckboxGroup.update(choices=[]), gr.Button.update(interactive=False)
+            return gr.update(value=f"❌ Init failed: {INIT_ERROR}"), gr.update(choices=[]), gr.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)
+        # Populate LoRA choices (manifest could come from repo, Space file, or built-in fallback)
+        return gr.update(value=msg), gr.update(choices=list(LORA_MANIFEST.keys())), gr.update(interactive=True)
 
     demo.load(_startup, outputs=[status, lora_names, btn])
 
@@ -246,5 +286,4 @@ 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()