Create app.py

app.py

@@ -0,0 +1,195 @@
+import os, io, json, base64
+from typing import List, Dict, Any, Optional
+from PIL import Image
+import torch
+import gradio as gr
+from huggingface_hub import snapshot_download
+from diffusers import (
+    StableDiffusionXLPipeline,
+    StableDiffusionPipeline,
+    DPMSolverMultistepScheduler,
+    EulerAncestralDiscreteScheduler,
+    EulerDiscreteScheduler,
+    DDIMScheduler,
+    LMSDiscreteScheduler,
+    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)
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+dtype = torch.float16 if device == "cuda" else torch.float32
+
+SCHEDULERS = {
+    "default": None,
+    "euler_a": EulerAncestralDiscreteScheduler,
+    "euler": EulerDiscreteScheduler,
+    "ddim": DDIMScheduler,
+    "lms": LMSDiscreteScheduler,
+    "pndm": PNDMScheduler,
+    "dpmpp_2m": DPMSolverMultistepScheduler,
+}
+
+# Globals populated at startup
+pipe = None
+IS_SDXL = True
+LORA_MANIFEST: Dict[str, Dict[str, str]] = {}
+
+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)
+    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:
+        pipe_local = StableDiffusionXLPipeline.from_single_file(
+            ckpt_path, torch_dtype=dtype, use_safetensors=True, add_watermarker=False
+        )
+        IS_SDXL = True
+    except Exception:
+        pipe_local = 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")
+    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 = {}
+
+    return pipe_local
+
+def apply_loras(selected: List[str], scale: float):
+    if not selected or scale <= 0:
+        return
+    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)
+            else:
+                pipe.load_lora_weights(meta.get("repo", ""), weight_name=meta.get("weight_name"), adapter_name=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))
+    except Exception as e:
+        print(f"[WARN] set_adapters failed: {e}")
+
+def txt2img(
+    prompt: str,
+    negative: str,
+    width: int,
+    height: int,
+    steps: int,
+    guidance: float,
+    images: int,
+    seed: Optional[int],
+    scheduler: str,
+    loras: List[str],
+    lora_scale: float,
+    fuse_lora: bool,
+):
+    if scheduler and 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(loras, lora_scale)
+    if fuse_lora and loras:
+        try:
+            pipe.fuse_lora(lora_scale=float(lora_scale))
+        except Exception as e:
+            print(f"[WARN] fuse_lora failed: {e}")
+
+    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,
+        width=int(width),
+        height=int(height),
+        num_inference_steps=int(steps),
+        guidance_scale=float(guidance),
+        num_images_per_prompt=int(images),
+        generator=generator,
+    )
+    out = pipe(**kwargs)
+    return out.images
+
+def warmup():
+    # Tiny, fast pass to initialize CUDA kernels and weight graphs
+    try:
+        _ = txt2img(
+            "warmup", "", 512 if IS_SDXL else 512, 512 if IS_SDXL else 512,
+            5, 5.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 (runs Diffusers directly)") as demo:
+    gr.Markdown("### SDXL text‑to‑image (single‑file checkpoint) with optional LoRAs")
+    with gr.Row():
+        prompt = gr.Textbox(label="Prompt", lines=3)
+        negative = gr.Textbox(label="Negative Prompt", lines=3)
+    with gr.Row():
+        width = gr.Slider(256, 1536, 1024, step=64, label="Width")
+        height = gr.Slider(256, 1536, 1024, step=64, label="Height")
+    with gr.Row():
+        steps = gr.Slider(5, 80, 30, step=1, label="Steps")
+        guidance = gr.Slider(0.0, 20.0, 6.5, step=0.1, label="Guidance")
+        images = gr.Slider(1, 4, 1, step=1, label="Images")
+    with gr.Row():
+        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_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)
+
+    def _startup():
+        global pipe
+        pipe = bootstrap_model()
+        # Fill LoRA choices after manifest loads
+        return gr.CheckboxGroup.update(choices=list(LORA_MANIFEST.keys()))
+    demo.load(_startup, outputs=[lora_names])
+    demo.load(lambda: warmup(), inputs=None, outputs=None)
+
+    btn.click(
+        txt2img,
+        inputs=[prompt, negative, width, height, steps, guidance, images, seed, scheduler, lora_names, lora_scale, fuse],
+        outputs=[gallery],
+        api_name="txt2img",
+    )
+
+# Tune queue for Spaces (concurrency + backlog)
+demo.queue(max_size=32, concurrency_count=2).launch()