app.py

"""gr.Workflow app for PrunaAI/P-Q-Image-2.1 (ZeroGPU).

Nodes on the canvas:
  - "Generate Image": text-to-image with a v0.2 or v0.1 LoRA adapter
  - "Edit Image": image editing with an input image + prompt
"""

import os
import spaces
import torch
import gradio as gr
from gradio.utils import get_upload_folder
from gradio_client import utils as client_utils
from PIL import Image
from diffusers import FlowMatchEulerDiscreteScheduler, QwenImage21Pipeline

LORA_REPO = "PrunaAI/P-Q-Image-2.1"
DEFAULT_VARIANT = "6-steps"
LORA_STRENGTH = 1.0

# Each adapter is trained for one schedule. v0.2 uses unshifted 6-step sigmas
# plus dynamic shifting; v0.1 uses its raw nodes with shifting off.
VARIANTS = {
    "6-steps": {
        "weight_name": "p_q_image_2.1_6_steps_v0.2.safetensors",
        "adapter_name": "step6_v02",
        "steps": 6,
        "sigmas": [1.0, 0.925, 0.85, 0.625, 0.375, 0.125],
        "dynamic_shifting": True,
        "use_kv_cache": False,
    },
    "6-steps-upsampled": {
        "weight_name": "p_q_image_2.1_6_steps_upsampled_v0.2.safetensors",
        "adapter_name": "step6_upsampled_v02",
        "steps": 6,
        "sigmas": [1.0, 0.925, 0.85, 0.625, 0.375, 0.125],
        "dynamic_shifting": True,
        "use_kv_cache": False,
    },
    "8-steps": {
        "weight_name": "p_q_image_2.1_8_steps_v0.1.safetensors",
        "adapter_name": "step8_v01",
        "steps": 8,
        "sigmas": [1.0, 14 / 15, 6 / 7, 10 / 13, 2 / 3, 6 / 11, 0.4, 2 / 9],
        "dynamic_shifting": False,
        "use_kv_cache": True,
    },
    "5-steps": {
        "weight_name": "p_q_image_2.1_5_steps_v0.1.safetensors",
        "adapter_name": "step5_v01",
        "steps": 5,
        "sigmas": [1.0, 0.94, 6 / 7, 2 / 3, 0.4],
        "dynamic_shifting": False,
        "use_kv_cache": True,
    },
}

VARIANT_CHOICES = list(VARIANTS)

VARIANT_ALIASES = {
    "6-steps": "6-steps",
    "6-step": "6-steps",
    "6step": "6-steps",
    "6": "6-steps",
    "basic": "6-steps",
    "v0.2": "6-steps",
    "v0.2-6-step": "6-steps",
    "v0.2-6-steps": "6-steps",
    "6-steps-upsampled": "6-steps-upsampled",
    "6-step-upsampled": "6-steps-upsampled",
    "6step-upsampled": "6-steps-upsampled",
    "prompt-upsampled": "6-steps-upsampled",
    "upsampled": "6-steps-upsampled",
    "upsample": "6-steps-upsampled",
    "v0.2-upsampled": "6-steps-upsampled",
    "v0.2-6-step-upsampled": "6-steps-upsampled",
    "v0.2-6-steps-upsampled": "6-steps-upsampled",
    "8-steps": "8-steps",
    "8-step": "8-steps",
    "8step": "8-steps",
    "8": "8-steps",
    "v0.1-8-step": "8-steps",
    "v0.1-8-steps": "8-steps",
    "5-steps": "5-steps",
    "5-step": "5-steps",
    "5step": "5-steps",
    "5": "5-steps",
    "v0.1-5-step": "5-steps",
    "v0.1-5-steps": "5-steps",
}

pipe = QwenImage21Pipeline.from_pretrained(
    "Qwen/Qwen-Image-2.1", torch_dtype=torch.bfloat16
)
for spec in VARIANTS.values():
    pipe.load_lora_weights(
        LORA_REPO,
        weight_name=spec["weight_name"],
        adapter_name=spec["adapter_name"],
    )
pipe.set_adapters([VARIANTS[DEFAULT_VARIANT]["adapter_name"]], [LORA_STRENGTH])
pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_config(
    pipe.scheduler.config,
    use_dynamic_shifting=True,
    shift_terminal=None,
)
pipe = pipe.to("cuda")

_active_variant = DEFAULT_VARIANT


def _resolve_variant(variant) -> str:
    if variant is None or variant == "":
        return DEFAULT_VARIANT
    key = "-".join(str(variant).strip().lower().replace("_", "-").split())
    try:
        key = str(int(float(key)))
    except (TypeError, ValueError):
        pass
    return VARIANT_ALIASES.get(key, DEFAULT_VARIANT)


def _ensure_variant(variant) -> dict:
    """Activate the matching adapter and scheduler. Load only one at a time."""
    global _active_variant
    name = _resolve_variant(variant)
    spec = VARIANTS[name]
    if _active_variant != name:
        pipe.set_adapters([spec["adapter_name"]], [LORA_STRENGTH])
        scheduler_kwargs = {"use_dynamic_shifting": spec["dynamic_shifting"], "shift_terminal": None}
        if not spec["dynamic_shifting"]:
            scheduler_kwargs["shift"] = 1.0
        pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_config(
            pipe.scheduler.config,
            **scheduler_kwargs,
        )
        _active_variant = name
    return spec


def _seed(seed) -> int:
    try:
        return int(seed)
    except (TypeError, ValueError):
        return 42


def _coerce_image(image):
    """Accept a PIL image, a filepath, or a Gradio/Workflow file-ref dict."""
    if image is None or image == "":
        return None
    if isinstance(image, Image.Image):
        return image.convert("RGB")
    if isinstance(image, dict):
        image = image.get("path") or image.get("url") or ""
    return Image.open(str(image)).convert("RGB")


def _to_file(img: Image.Image) -> dict:
    """Save a PIL image into Gradio's upload folder and return a file-ref dict
    (same shape gr.Workflow uses internally for media node outputs)."""
    directory = get_upload_folder()
    os.makedirs(directory, exist_ok=True)
    path = os.path.join(directory, f"pruna_{os.urandom(8).hex()}.png")
    img.save(path)
    url = f"/gradio_api/file={client_utils.encode_file_path(path)}"
    return {"path": path, "url": url, "is_file": True}


def _run_pipe(prompt: str, spec: dict, seed: int, image=None):
    kwargs = {
        "prompt": prompt,
        "generator": torch.Generator("cuda").manual_seed(_seed(seed)),
        "num_inference_steps": spec["steps"],
        "sigmas": spec["sigmas"],
        "true_cfg_scale": 1.0,
        "use_kv_cache": spec["use_kv_cache"],
        "output_resolution": 1024,
    }
    if image is None:
        kwargs["width"] = 1024
        kwargs["height"] = 1024
    else:
        kwargs["image"] = image
    return pipe(**kwargs).images[0]


def _variant_dropdown() -> gr.Dropdown:
    return gr.Dropdown(
        choices=VARIANT_CHOICES,
        value=DEFAULT_VARIANT,
        label="Variant",
    )


@spaces.GPU(duration=120)
def generate(
    prompt: str,
    variant: str = _variant_dropdown(),
    seed: int = 42,
) -> dict:
    """Text-to-image with P-Q-Image-2.1 (no CFG)."""
    spec = _ensure_variant(variant)
    return _to_file(_run_pipe(prompt, spec, seed))


@spaces.GPU(duration=120)
def edit(
    image,
    prompt: str,
    variant: str = _variant_dropdown(),
    seed: int = 42,
) -> dict:
    """Edit an image with P-Q-Image-2.1 (no CFG)."""
    spec = _ensure_variant(variant)
    return _to_file(_run_pipe(prompt, spec, seed, image=_coerce_image(image)))


demo = gr.Workflow(
    graph="workflow.json",
    bind={"Generate Image": generate, "Edit Image": edit},
)

if __name__ == "__main__":
    demo.launch()