app.py

import gc, random
import gradio as gr
import torch, spaces
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler

# --- gradio_client bool-schema hotfix (safe no-op if not needed) ---
try:
    import gradio_client.utils as _gcu
    _orig_get_type = _gcu.get_type
    def _get_type_safe(schema):
        # Handle JSON Schema booleans (e.g., additionalProperties: True/False)
        if isinstance(schema, bool):
            return "any"
        return _orig_get_type(schema)
    _gcu.get_type = _get_type_safe
except Exception:
    pass
# -------------------------------------------------------------------

# sanity: show CPU at startup (no GPU yet)
zero = torch.tensor([0.0])
print("startup device:", zero.device)  # should print: cpu

# ====== 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)

@spaces.GPU  # tiny demo proving CUDA is present inside GPU block
def greet(n: float):
    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()

    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]

# ====== UI ======
with gr.Blocks() as demo:
    gr.Markdown("# ZeroGPU demo + Stable Diffusion 1.5 (minimal)")

    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)

    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__":
    # On Spaces, expose a shareable URL & bind to 0.0.0.0
    demo.queue(max_size=12).launch(share=False, show_error=True, mcp_server=True)