Test ZeroGPU

app.py

@@ -1,30 +1,99 @@
+import gc, random
 import gradio as gr
-import spaces
-import torch
+import torch, spaces
+from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
 
-# Outside of GPU context → no GPU available
-zero = torch.tensor([0.0])  # stays on CPU at import time
-print("startup device:", zero.device)  # will log: cpu
+# ----- sanity: show CPU at startup (no GPU yet) -----
+zero = torch.tensor([0.0])
+print("startup device:", zero.device)  # should print: cpu
 
-@spaces.GPU  # GPU is allocated ONLY while this function runs
+# ====== SD 1.5 minimal loader (lazy) ======
+MODEL_ID = "runwayml/stable-diffusion-v1-5"
+DTYPE = torch.float16
+_PIPE = None
+
+def get_pipe():
+    """Create the pipeline on first use (safe for ZeroGPU)."""
+    global _PIPE
+    if _PIPE is None:
+        pipe = StableDiffusionPipeline.from_pretrained(
+            MODEL_ID,
+            torch_dtype=DTYPE,
+            safety_checker=None,
+            use_safetensors=True,
+            low_cpu_mem_usage=True,
+        )
+        # fast/stable scheduler
+        pipe.scheduler = DPMSolverMultistepScheduler.from_config(
+            pipe.scheduler.config, use_karras_sigmas=True, algorithm_type="dpmsolver++"
+        )
+        # memory savers
+        pipe.enable_attention_slicing()
+        pipe.enable_vae_slicing()
+        pipe.enable_model_cpu_offload()
+        _PIPE = pipe
+    return _PIPE
+
+def snap8(x: int) -> int:
+    x = max(256, min(1024, int(x)))
+    return x - (x % 8)
+
+# ====== original ZeroGPU example ======
+@spaces.GPU
 def greet(n: float):
-    # Inside GPU context → CUDA is available
-    print("inside greet, torch.cuda.is_available():", torch.cuda.is_available())
+    print("inside greet, cuda available:", torch.cuda.is_available())
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    t = zero.to(device) + torch.tensor([float(n)], device=device)
+    return f"Hello {t.item():.3f} Tensor (device: {t.device})"
+
+# ====== SD 1.5 text -> image ======
+@spaces.GPU(duration=120)
+def generate(prompt: str, negative: str, steps: int, cfg: float, width: int, height: int, seed: int):
+    pipe = get_pipe()
+    w, h = snap8(width), snap8(height)
+
+    # seed
+    if int(seed) < 0:
+        seed = random.randint(0, 2**31 - 1)
+    gen = torch.Generator(device="cuda").manual_seed(int(seed))
+
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+    gc.collect()
 
-    # Move the tensor to the active device and add the input
-    t0 = zero.to(device)
-    tn = torch.tensor([float(n)], device=device)
-    out = t0 + tn
+    with torch.autocast(device_type="cuda", dtype=DTYPE):
+        out = pipe(
+            prompt=str(prompt),
+            negative_prompt=str(negative or ""),
+            num_inference_steps=int(steps),
+            guidance_scale=float(cfg),
+            width=w, height=h,
+            generator=gen,
+        )
+    return out.images[0]
 
-    # Log where it ran; on ZeroGPU this should be cuda:0
-    print("tensor device during compute:", out.device)
+# ====== UI (two tabs: sanity + SD) ======
+with gr.Blocks() as demo:
+    gr.Markdown("# ZeroGPU demo + Stable Diffusion 1.5 (minimal)")
 
-    # Return a friendly message
-    return f"Hello {out.item():.3f} (device: {out.device})"
+    with gr.Tab("ZeroGPU sanity"):
+        n = gr.Number(label="Add to zero", value=1.0)
+        hello = gr.Textbox(label="Result")
+        gr.Button("Greet").click(greet, n, hello)
 
-demo = gr.Interface(fn=greet, inputs=gr.Number(label="Add to zero"), outputs=gr.Text(label="Result"))
+    with gr.Tab("SD 1.5: Text → Image"):
+        prompt  = gr.Textbox(label="Prompt", value="a cozy reading nook, warm sunlight, cinematic lighting, highly detailed")
+        negative = gr.Textbox(label="Negative (optional)", value="lowres, blurry, watermark, text")
+        steps   = gr.Slider(8, 40, value=28, step=1, label="Steps")
+        cfg     = gr.Slider(1.0, 12.0, value=7.0, step=0.5, label="CFG")
+        width   = gr.Slider(256, 1024, value=640, step=16, label="Width")
+        height  = gr.Slider(256, 1024, value=640, step=16, label="Height")
+        seed    = gr.Number(value=-1, precision=0, label="Seed (-1 random)")
+        out_img = gr.Image(label="Image", interactive=False)
+        gr.Button("Generate").click(
+            generate, [prompt, negative, steps, cfg, width, height, seed], out_img
+        )
 
 if __name__ == "__main__":
-    # Minimal launch so the Space just builds
+    # keep it minimal so the Space "just builds"
     demo.launch()