app.py

import spaces  # must be imported before torch / diffusers

import gc
import random
from urllib.parse import unquote

import gradio as gr
import numpy as np
import torch
from PIL import Image
from diffusers import AutoPipelineForImage2Image, LCMScheduler
from safetensors.torch import load_file as load_safetensors

MAX_SEED = np.iinfo(np.int32).max
device = "cuda"
dtype = torch.float16

DEFAULT_BASE = "stable-diffusion-v1-5/stable-diffusion-v1-5"
LCM_LORA = "latent-consistency/lcm-lora-sdv1-5"
FAST_CHOICES = ["Fast (LCM, ~6 steps)", "Normal (30 steps)"]

# what is currently loaded
STATE = {"base": None, "pipe": None, "scheduler": None, "lora": None, "fast": None}


def get_pipe(base_model):
    base_model = (base_model or DEFAULT_BASE).strip() or DEFAULT_BASE
    if STATE["base"] != base_model:
        if STATE["pipe"] is not None:
            STATE["pipe"] = None
            gc.collect()
            torch.cuda.empty_cache()
        pipe = AutoPipelineForImage2Image.from_pretrained(
            base_model,
            torch_dtype=dtype,
            safety_checker=None,
            requires_safety_checker=False,
        ).to(device)
        pipe.set_progress_bar_config(disable=True)
        STATE.update(base=base_model, pipe=pipe,
                     scheduler=pipe.scheduler, lora=None, fast=None)
    return STATE["pipe"]


get_pipe(DEFAULT_BASE)  # warm up at startup


def parse_lora_ref(ref, fname):
    """Accepts 'user/repo', a full HF link, or a link to a .safetensors file."""
    ref = (ref or "").strip()
    fname = (fname or "").strip()
    if not ref:
        return None, None
    if "huggingface.co" in ref:
        part = unquote(ref.split("huggingface.co/", 1)[1]).split("?")[0].strip("/")
        segs = [s for s in part.split("/") if s]
        if segs and segs[0] in ("models", "spaces", "datasets"):
            segs = segs[1:]
        repo = "/".join(segs[:2])
        if len(segs) > 4 and segs[2] in ("blob", "resolve"):
            fname = fname or "/".join(segs[4:])
        return repo, (fname or None)
    return ref, (fname or None)


def apply_adapters(pipe, fast_mode, lora_repo, lora_file, lora_path):
    fast = fast_mode.startswith("Fast")
    lora_key = lora_path or ((lora_repo, lora_file) if lora_repo else None)

    if (fast, lora_key) != (STATE["fast"], STATE["lora"]):
        try:
            pipe.unload_lora_weights()
        except Exception:
            pass
        STATE["fast"], STATE["lora"] = None, None

        names = []
        if fast:
            pipe.load_lora_weights(LCM_LORA, adapter_name="lcm")
            names.append("lcm")
        if lora_key:
            if lora_path:
                pipe.load_lora_weights(load_safetensors(lora_path), adapter_name="user")
            elif lora_file:
                pipe.load_lora_weights(lora_repo, weight_name=lora_file, adapter_name="user")
            else:
                pipe.load_lora_weights(lora_repo, adapter_name="user")
            names.append("user")
        STATE["fast"], STATE["lora"] = fast, lora_key

    pipe.scheduler = (LCMScheduler.from_config(STATE["scheduler"].config)
                      if fast else STATE["scheduler"])
    return ["lcm"] * int(fast) + (["user"] if lora_key else [])


@spaces.GPU(duration=90)
def infer(
    image,
    prompt,
    negative_prompt,
    base_model,
    fast_mode,
    lora_ref,
    lora_weight_name,
    lora_upload,
    lora_scale,
    strength,
    steps,
    guidance_scale,
    size,
    seed,
    randomize_seed,
    progress=gr.Progress(track_tqdm=True),
):
    if image is None:
        raise gr.Error("Please upload an image.")
    if not (prompt or "").strip():
        raise gr.Error("Please enter a prompt.")

    try:
        pipe = get_pipe(base_model)
    except Exception as e:
        STATE["base"] = None
        raise gr.Error(f"Could not load base model: {e}")

    lora_repo, lora_file = parse_lora_ref(lora_ref, lora_weight_name)
    lora_path = lora_upload if isinstance(lora_upload, str) else getattr(lora_upload, "name", None)

    try:
        names = apply_adapters(pipe, fast_mode, lora_repo, lora_file, lora_path)
    except Exception as e:
        STATE["fast"], STATE["lora"] = None, None
        raise gr.Error(f"Could not load LoRA: {e}")

    if names:
        pipe.set_adapters(names, adapter_weights=[
            1.0 if n == "lcm" else float(lora_scale) for n in names
        ])

    img = image.convert("RGB")
    long_side = int(size)
    w, h = img.size
    if w >= h:
        nw, nh = long_side, int(long_side * h / w)
    else:
        nh, nw = long_side, int(long_side * w / h)
    img = img.resize(((nw // 8) * 8, (nh // 8) * 8), Image.LANCZOS)

    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    generator = torch.Generator(device=device).manual_seed(int(seed))

    try:
        out = pipe(
            prompt=prompt,
            negative_prompt=negative_prompt or None,
            image=img,
            strength=float(strength),
            num_inference_steps=int(steps),
            guidance_scale=float(guidance_scale),
            generator=generator,
        ).images[0]
        return out, int(seed)
    finally:
        gc.collect()
        torch.cuda.empty_cache()


def on_fast_change(fast_mode):
    if fast_mode.startswith("Fast"):
        return gr.update(value=6), gr.update(value=1.5)
    return gr.update(value=30), gr.update(value=7.5)


css = """
.gradio-container{max-width:1300px!important}
footer{display:none!important}
#run-btn{font-size:16px;font-weight:600}
#speed-radio label{border:1px solid var(--border-color-primary);border-radius:8px;
  padding:8px 14px;margin:4px 6px 4px 0;font-weight:600;cursor:pointer}
#speed-radio label:has(input:checked){border-color:var(--color-accent);
  background:var(--color-accent-soft)}
"""

with gr.Blocks() as demo:
    gr.Markdown("## Stable Diffusion 1.5 — image to image, bring your own LoRA")

    with gr.Row():
        with gr.Column(scale=1):
            image = gr.Image(label="Input image", type="pil", height=380)
            prompt = gr.Textbox(label="Prompt", lines=3,
                                placeholder="e.g. oil painting, dramatic lighting, highly detailed")
            strength = gr.Slider(0.1, 1.0, value=0.6, step=0.05, label="Strength",
                                 info="How much to change the original. 0.3 = light touch, 0.8 = almost a new image.")
            with gr.Row():
                run_btn = gr.Button("Run", variant="primary", elem_id="run-btn")
                stop_btn = gr.Button("Stop")

        with gr.Column(scale=1):
            result = gr.Image(label="Result", format="png", height=460)

            fast_mode = gr.Radio(
                FAST_CHOICES, value="Fast (LCM, ~6 steps)", label="Speed",
                elem_id="speed-radio",
                info="Fast uses the LCM LoRA. Normal uses the plain model: slower, usually better.",
            )

            with gr.Accordion("My LoRA", open=True):
                lora_ref = gr.Textbox(label="HF repo or link",
                                      placeholder="myuser/my-lora   or   https://huggingface.co/.../file.safetensors")
                lora_weight_name = gr.Textbox(label="File name (optional)",
                                              placeholder="my_lora.safetensors")
                lora_upload = gr.File(label="...or upload a .safetensors from your computer",
                                      file_types=[".safetensors"], type="filepath")
                lora_scale = gr.Slider(0.0, 2.0, value=1.0, step=0.05, label="LoRA strength")

            with gr.Accordion("Settings", open=False):
                base_model = gr.Textbox(label="Base model", value=DEFAULT_BASE,
                                        info="Any SD 1.5 checkpoint in diffusers format, e.g. Lykon/dreamshaper-8")
                size = gr.Radio([512, 640, 768], value=640, label="Output size (long side)")
                steps = gr.Slider(1, 50, value=6, step=1, label="Steps")
                guidance_scale = gr.Slider(0.0, 15.0, value=1.5, step=0.1, label="Guidance (CFG)")
                negative_prompt = gr.Textbox(label="Negative prompt",
                                             value="worst quality, low quality, blurry, bad anatomy, bad hands, watermark, text")
                seed = gr.Slider(0, MAX_SEED, value=0, step=1, label="Seed")
                randomize_seed = gr.Checkbox(value=True, label="Random seed")

    fast_mode.change(on_fast_change, inputs=fast_mode, outputs=[steps, guidance_scale])

    args = [image, prompt, negative_prompt, base_model, fast_mode, lora_ref,
            lora_weight_name, lora_upload, lora_scale, strength, steps,
            guidance_scale, size, seed, randomize_seed]

    ev = run_btn.click(fn=infer, inputs=args, outputs=[result, seed])
    prompt.submit(fn=infer, inputs=args, outputs=[result, seed])
    stop_btn.click(fn=None, inputs=None, outputs=None, cancels=[ev])

if __name__ == "__main__":
    demo.queue(max_size=20).launch(
        css=css,
        theme=gr.themes.Soft(),
        ssr_mode=False,
        show_error=True,
    )