worker/handler.py

"""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})