"""RunPod serverless worker: Qwen-Image + Lightning (turbo) with per-request LoRA URLs. Input schema (all inside "input"): prompt str, required negative_prompt str, default " " width / height int, default 1328x1328 (or "size": "1344*768") num_inference_steps int, default 8 (lightning) true_cfg_scale float, default 1.0 (lightning; use 4.0 + ~50 steps without lightning) seed int, default random num_images int, default 1 (max 4) loras list of {"url": str, "scale": float} — downloaded and applied per request (also accepts lora_url/lora_scale shorthand) output_format "png" | "jpeg", default "png" Returns: {"images": [base64...], "seed": int, "timings": {...}} No safety checker / content filter is present in this pipeline. """ import base64 import hashlib import io import math import os import time import traceback import urllib.request import torch from safetensors.torch import load_file MODEL_ID = os.environ.get("MODEL_ID", "Qwen/Qwen-Image") LIGHTNING_REPO = os.environ.get("LIGHTNING_REPO", "lightx2v/Qwen-Image-Lightning") LIGHTNING_FILE = os.environ.get( "LIGHTNING_FILE", "Qwen-Image-Lightning-8steps-V2.0-bf16.safetensors" ) HF_TOKEN = os.environ.get("HF_TOKEN") LORA_CACHE = "/lora-cache" os.makedirs(LORA_CACHE, exist_ok=True) # Scheduler config recommended by lightx2v for Lightning checkpoints. LIGHTNING_SCHEDULER = { "base_image_seq_len": 256, "base_shift": math.log(3), "invert_sigmas": False, "max_image_seq_len": 8192, "max_shift": math.log(3), "num_train_timesteps": 1000, "shift": 1.0, "shift_terminal": None, "stochastic_sampling": False, "time_shift_type": "exponential", "use_beta_sigmas": False, "use_dynamic_shifting": True, "use_exponential_sigmas": False, "use_karras_sigmas": False, } print(f"[init] loading {MODEL_ID} ...", flush=True) t0 = time.time() from diffusers import DiffusionPipeline, FlowMatchEulerDiscreteScheduler # noqa: E402 scheduler = FlowMatchEulerDiscreteScheduler.from_config(LIGHTNING_SCHEDULER) pipe = DiffusionPipeline.from_pretrained( MODEL_ID, scheduler=scheduler, torch_dtype=torch.bfloat16, token=HF_TOKEN ) pipe.to("cuda") print(f"[init] pipeline loaded in {time.time()-t0:.0f}s", flush=True) if LIGHTNING_FILE.lower() not in ("", "none", "off"): t1 = time.time() pipe.load_lora_weights( LIGHTNING_REPO, weight_name=LIGHTNING_FILE, adapter_name="lightning", token=HF_TOKEN ) pipe.fuse_lora() pipe.unload_lora_weights() print(f"[init] lightning fused in {time.time()-t1:.0f}s", flush=True) def _download(url: str) -> str: path = os.path.join(LORA_CACHE, hashlib.sha1(url.encode()).hexdigest() + ".safetensors") if os.path.exists(path): return path headers = {} if HF_TOKEN and "huggingface.co" in url: headers["Authorization"] = f"Bearer {HF_TOKEN}" req = urllib.request.Request(url, headers=headers) tmp = path + ".part" with urllib.request.urlopen(req, timeout=300) as r, open(tmp, "wb") as f: while chunk := r.read(1 << 20): f.write(chunk) os.rename(tmp, path) return path def _load_lora_state(path: str) -> dict: sd = load_file(path) out = {} for k, v in sd.items(): if v.dtype in (torch.float8_e4m3fn, torch.float8_e5m2): v = v.to(torch.bfloat16) # ai-toolkit / comfy prefix -> diffusers prefix if k.startswith("diffusion_model."): k = "transformer." + k[len("diffusion_model."):] out[k] = v return out def handler(job): inp = job.get("input") or {} prompt = inp.get("prompt") if not prompt: return {"error": "input.prompt is required"} if "size" in inp: try: w, h = (int(x) for x in str(inp["size"]).replace("x", "*").split("*")) except Exception: return {"error": f"bad size: {inp['size']}"} else: w, h = int(inp.get("width", 1328)), int(inp.get("height", 1328)) w, h = max(64, w - w % 16), max(64, h - h % 16) steps = int(inp.get("num_inference_steps", inp.get("steps", 8))) cfg = float(inp.get("true_cfg_scale", inp.get("cfg", inp.get("guidance", 1.0)))) num_images = min(int(inp.get("num_images", 1)), 4) seed = inp.get("seed") if seed is None or int(seed) < 0: seed = torch.seed() % (2**31) seed = int(seed) loras = list(inp.get("loras") or []) if inp.get("lora_url"): loras.append({"url": inp["lora_url"], "scale": inp.get("lora_scale", 1.0)}) timings = {} adapters, scales = [], [] try: t = time.time() for i, l in enumerate(loras): url = l.get("url") or l.get("path") if not url: return {"error": f"loras[{i}] needs url"} name = f"user{i}" pipe.load_lora_weights(_load_lora_state(_download(url)), adapter_name=name) adapters.append(name) scales.append(float(l.get("scale", 1.0))) if adapters: pipe.set_adapters(adapters, adapter_weights=scales) timings["lora_s"] = round(time.time() - t, 1) t = time.time() gen = torch.Generator(device="cuda").manual_seed(seed) images = pipe( prompt=prompt, negative_prompt=inp.get("negative_prompt", " "), width=w, height=h, num_inference_steps=steps, true_cfg_scale=cfg, num_images_per_prompt=num_images, generator=gen, ).images timings["generate_s"] = round(time.time() - t, 1) fmt = str(inp.get("output_format") or inp.get("image_format") or "png").lower() fmt = {"jpg": "JPEG", "jpeg": "JPEG", "webp": "WEBP"}.get(fmt, "PNG") quality = int(inp.get("image_quality", 95)) out = [] for img in images: buf = io.BytesIO() img.save(buf, format=fmt, quality=quality) out.append(base64.b64encode(buf.getvalue()).decode()) return {"images": out, "seed": seed, "width": w, "height": h, "timings": timings} except Exception as e: traceback.print_exc() return {"error": f"{type(e).__name__}: {e}"} finally: if adapters: try: pipe.unload_lora_weights() except Exception: traceback.print_exc() import runpod # noqa: E402 runpod.serverless.start({"handler": handler})