Update app.py

app.py

@@ -6,8 +6,8 @@ import gradio as gr
 import spaces
 from huggingface_hub import snapshot_download
 from diffusers import (
-    StableDiffusionXLPipeline,
-    StableDiffusionPipeline,
+    StableDiffusionPipeline,                 # SD 1.x/2.x single-file loader
+    StableDiffusionXLPipeline,               # SDXL single-file loader
     DPMSolverMultistepScheduler,
     EulerAncestralDiscreteScheduler,
     EulerDiscreteScheduler,
@@ -16,16 +16,13 @@ from diffusers import (
     PNDMScheduler,
 )
 
-# ----------------- Config (set in Space Secrets if private) -----------------
-MODEL_REPO_ID = os.getenv("MODEL_REPO_ID", "DB2169/Mixy").strip()
+# -------- Config --------
+MODEL_REPO_ID = os.getenv("MODEL_REPO_ID", "DB2169/mixy").strip()
 CHECKPOINT_FILENAME = os.getenv("CHECKPOINT_FILENAME", "unknown.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
+DO_WARMUP = os.getenv("WARMUP", "1") == "1"
 
-# 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 = {
@@ -38,22 +35,14 @@ SCHEDULERS = {
     "dpmpp_2m": DPMSolverMultistepScheduler,
 }
 
-# Globals populated at startup
+# -------- Globals --------
 pipe = None
-IS_SDXL = True
+IS_SDXL = False
 LORA_MANIFEST: Dict[str, Dict[str, str]] = {}
 INIT_ERROR: Optional[str] = None
 
-# ----------------- Helpers -----------------
+# -------- 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)
@@ -62,7 +51,6 @@ def load_lora_manifest(repo_dir: str) -> Dict[str, Dict[str, str]]:
         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:
@@ -73,7 +61,6 @@ def load_lora_manifest(repo_dir: str) -> Dict[str, Dict[str, str]]:
         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:
@@ -84,7 +71,6 @@ def load_lora_manifest(repo_dir: str) -> Dict[str, Dict[str, str]]:
         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": {
@@ -93,11 +79,11 @@ def load_lora_manifest(repo_dir: str) -> Dict[str, Dict[str, str]]:
         }
     }
 
-# ----------------- Bootstrap (download + load on CPU) -----------------
+# -------- Bootstrap (CPU) --------
 def bootstrap_model():
     """
-    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.
+    Try SD (1.x/2.x) single-file first, then SDXL single-file, to maximize compatibility
+    with older diffusers that don’t expose DiffusionPipeline.from_single_file.
     """
     global pipe, IS_SDXL, LORA_MANIFEST, INIT_ERROR
     INIT_ERROR = None
@@ -125,20 +111,24 @@ def bootstrap_model():
         print(f"[ERROR] {INIT_ERROR}")
         return
 
+    _pipe = None
+    _is_sdxl = False
+
+    # 1) SD 1.x/2.x first (most single-file merges are SD), then SDXL
     try:
-        # Attempt SDXL first (text_encoder_2 present)
-        _pipe = StableDiffusionXLPipeline.from_single_file(
-            ckpt_path, torch_dtype=torch.float16, use_safetensors=True, add_watermarker=False
+        _pipe = StableDiffusionPipeline.from_single_file(
+            ckpt_path, torch_dtype=torch.float16, use_safetensors=True
         )
-        sdxl = True
-    except Exception:
+        _is_sdxl = False
+    except Exception as e_sd:
+        print(f"[INFO] SD load failed or not SD: {e_sd}")
         try:
-            _pipe = StableDiffusionPipeline.from_single_file(
-                ckpt_path, torch_dtype=torch.float16, use_safetensors=True
+            _pipe = StableDiffusionXLPipeline.from_single_file(
+                ckpt_path, torch_dtype=torch.float16, use_safetensors=True, add_watermarker=False
             )
-            sdxl = False
-        except Exception as e:
-            INIT_ERROR = f"Failed to load pipeline: {e}"
+            _is_sdxl = True
+        except Exception as e_sdxl:
+            INIT_ERROR = f"Failed to load pipeline (SD and SDXL): SD={e_sd} | SDXL={e_sdxl}"
             print(f"[ERROR] {INIT_ERROR}")
             return
 
@@ -152,9 +142,8 @@ def bootstrap_model():
     manifest = load_lora_manifest(local_dir)
     print(f"[INFO] LoRAs available: {list(manifest.keys())}")
 
-    # Publish
     pipe = _pipe
-    IS_SDXL = sdxl
+    IS_SDXL = _is_sdxl
     LORA_MANIFEST = manifest
 
 def apply_loras(selected: List[str], scale: float, repo_dir: str):
@@ -179,7 +168,7 @@ def apply_loras(selected: List[str], scale: float, repo_dir: str):
     except Exception as e:
         print(f"[WARN] set_adapters failed: {e}")
 
-# ----------------- Generation (ZeroGPU) -----------------
+# -------- Generation (ZeroGPU) --------
 @spaces.GPU
 def txt2img(
     prompt: str,
@@ -230,14 +219,8 @@ 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}")
-
-# ----------------- UI -----------------
-with gr.Blocks(title="SDXL Space (ZeroGPU, single-file, LoRA-ready)") as demo:
+# -------- UI --------
+with gr.Blocks(title="SDXL/SD single-file (ZeroGPU, LoRA-ready)") as demo:
     status = gr.Markdown("")
 
     with gr.Row():
@@ -267,15 +250,30 @@ with gr.Blocks(title="SDXL Space (ZeroGPU, single-file, LoRA-ready)") as demo:
     def _startup():
         bootstrap_model()
         if INIT_ERROR:
-            return gr.update(value=f"❌ Init failed: {INIT_ERROR}"), gr.update(choices=[]), gr.update(interactive=False)
+            return (
+                gr.update(value=f"❌ Init failed: {INIT_ERROR}"),
+                gr.update(choices=[]),
+                gr.update(value=1024, minimum=256, maximum=1536, step=64),
+                gr.update(value=1024, minimum=256, maximum=1536, step=64),
+                gr.update(interactive=False),
+            )
+        default_wh = 1024 if IS_SDXL else 512
         msg = f"✅ Model loaded from {MODEL_REPO_ID} ({'SDXL' if IS_SDXL else 'SD'})"
-        # 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])
+        # Warm up only after model is ready (avoids race)
+        if DO_WARMUP:
+            try:
+                _ = txt2img("warmup", "", default_wh, default_wh, 4, 4.0, 1, 1234, "default", [], 0.0, False)
+            except Exception as e:
+                print(f"[WARN] Warmup failed: {e}")
+        return (
+            gr.update(value=msg),
+            gr.update(choices=list(LORA_MANIFEST.keys())),
+            gr.update(value=default_wh, minimum=256, maximum=1536, step=64),
+            gr.update(value=default_wh, minimum=256, maximum=1536, step=64),
+            gr.update(interactive=True),
+        )
 
-    if DO_WARMUP:
-        demo.load(lambda: warmup(), inputs=None, outputs=None)
+    demo.load(_startup, outputs=[status, lora_names, width, height, btn])
 
     btn.click(
         txt2img,