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, )