Blame for app.py

"""P-Q-Image-2.1-Remove-Object Studio — gr.Server studio + gr.Workflow canvas.

  - "/"          Studio: custom UI (studio.html) over the `remove` API endpoint
  - "/workflow"  Workflow canvas: node graph from workflow.json (pq_remove)

The removal LoRA runs as its model card says: the few-step variants on top of PrunaAI's distillation
LoRA (at 1.0) at strength 0.7, the base variant alone at 40 steps; no CFG in all cases.
"""

import base64
import io
import json
import os
import random
import time
from urllib.parse import unquote, urlparse

import gradio as gr
import numpy as np
import spaces
import torch
import uvicorn
from fastapi import HTTPException
from fastapi.responses import FileResponse, HTMLResponse
from gradio.utils import get_upload_folder
from gradio_client import utils as client_utils
from PIL import Image, ImageOps
from diffusers import FlowMatchEulerDiscreteScheduler, QwenImage21Pipeline

from draw_outline import draw_outline, outline_prompt
from examples import EXAMPLES

APP_DIR = os.path.dirname(os.path.abspath(__file__))
EXAMPLES_DIR = os.path.join(APP_DIR, "examples")
MAX_SEED = 2**31 - 1

REMOVAL_REPO = "PrunaAI/P-Q-Image-2.1-Remove-Object"
DISTILL_REPO = "PrunaAI/P-Q-Image-2.1"
DEFAULT_VARIANT = "6-step"
DEFAULT_INSTRUCTION = "Remove the outlined object."

# From the model card's Variants table. `max_side` caps the working resolution (side of the
# equivalent square): ~2 MP like the card's examples for the few-step variants, ~1 MP for the
# 40-step base variant to keep it within a ZeroGPU slot. `gpu_seconds` is the ZeroGPU reservation.
VARIANTS = {
    "6-step": {
        "removal_weight": "p_q_image_2.1_remove_object_6_steps_v0.1.safetensors",
        "distill_weight": "p_q_image_2.1_6_steps_v0.2.safetensors",
        "strength": 0.7,
        # Unshifted sigmas: the scheduler shifts them by image size.
        "sigmas": [1.0, 0.925, 0.85, 0.625, 0.375, 0.125],
        "dynamic_shifting": True,
        "max_side": 1414,
        "gpu_seconds": 60,
    },
    "8-step": {
        "removal_weight": "p_q_image_2.1_remove_object_8_steps_v0.1.safetensors",
        "distill_weight": "p_q_image_2.1_8_steps_v0.1.safetensors",
        "strength": 0.7,
        # Sigmas used exactly as given.
        "sigmas": [1.0, 14 / 15, 6 / 7, 10 / 13, 2 / 3, 6 / 11, 0.4, 2 / 9],
        "dynamic_shifting": False,
        "max_side": 1414,
        "gpu_seconds": 60,
    },
    "base": {
        "removal_weight": "p_q_image_2.1_remove_object_base_v0.1.safetensors",
        "distill_weight": None,
        "strength": 1.0,
        "sigmas": None,  # 40 steps, the pipeline's own scheduler
        "dynamic_shifting": None,
        "max_side": 1024,
        "gpu_seconds": 150,
    },
}
VARIANT_CHOICES = list(VARIANTS)
VARIANT_ALIASES = {
    "6-step": "6-step", "6-steps": "6-step", "6step": "6-step", "6": "6-step",
    "8-step": "8-step", "8-steps": "8-step", "8step": "8-step", "8": "8-step",
    "base": "base", "40-step": "base", "40-steps": "base", "40": "base",
}

# ------------------------------------------------------------------ pipeline

pipe = QwenImage21Pipeline.from_pretrained("Qwen/Qwen-Image-2.1", torch_dtype=torch.bfloat16)
for name, spec in VARIANTS.items():
    spec["removal_adapter"] = f"removal_{name}"
    pipe.load_lora_weights(REMOVAL_REPO, weight_name=spec["removal_weight"], adapter_name=spec["removal_adapter"])
    if spec["distill_weight"]:
        spec["distill_adapter"] = f"distill_{name}"
        pipe.load_lora_weights(DISTILL_REPO, weight_name=spec["distill_weight"], adapter_name=spec["distill_adapter"])

# One scheduler per variant, built once: the pipeline's own for base, and for the few-step variants
# the schedule the distillation LoRA was trained with (no terminal stretch).
SCHEDULERS = {"base": pipe.scheduler}
for name, spec in VARIANTS.items():
    if spec["distill_weight"]:
        SCHEDULERS[name] = FlowMatchEulerDiscreteScheduler.from_config(
            pipe.scheduler.config, use_dynamic_shifting=spec["dynamic_shifting"], shift=1.0, shift_terminal=None
        )

_active_variant = None


def _activate(name: str) -> dict:
    """Switch adapters and scheduler to `name`; a no-op when it is already active."""
    global _active_variant
    spec = VARIANTS[name]
    if _active_variant != name:
        if spec["distill_weight"]:
            pipe.set_adapters([spec["distill_adapter"], spec["removal_adapter"]], [1.0, spec["strength"]])
        else:
            pipe.set_adapters([spec["removal_adapter"]], [spec["strength"]])
        pipe.scheduler = SCHEDULERS[name]
        _active_variant = name
    return spec


_activate(DEFAULT_VARIANT)
pipe = pipe.to("cuda")


def _resolve_variant(variant) -> str:
    if variant is None or variant == "":
        return DEFAULT_VARIANT
    key = "-".join(str(variant).strip().lower().replace("_", "-").split())
    if key not in VARIANT_ALIASES:
        raise gr.Error(f"Variant must be one of: {', '.join(VARIANT_CHOICES)}.")
    return VARIANT_ALIASES[key]


# ------------------------------------------------------------------ rendering


def _gpu_seconds(prompt, variant, seed, image, resolution):
    return VARIANTS[variant]["gpu_seconds"]


@spaces.GPU(duration=_gpu_seconds)
def _render(prompt: str, variant: str, seed: int, image: Image.Image, resolution: int):
    spec = _activate(variant)
    torch.cuda.synchronize()
    start = time.perf_counter()
    result = pipe(
        prompt=prompt,
        image=[image],
        sigmas=spec["sigmas"],
        num_inference_steps=len(spec["sigmas"]) if spec["sigmas"] else 40,
        true_cfg_scale=1.0,
        output_resolution=resolution,
        generator=torch.Generator("cuda").manual_seed(seed),
    ).images[0]
    torch.cuda.synchronize()
    return result, time.perf_counter() - start


def _canvas(ratio: float, side: int) -> tuple[int, int]:
    """The pipeline's own sizing: area side², this aspect ratio, sides rounded to multiples of 32."""
    width = (side**2 * ratio) ** 0.5
    return round(width / 32) * 32, round(width / ratio / 32) * 32


def _remove(photo: Image.Image, mask, instruction: str, variant, seed: int):
    """Object removal as the model card does it: work at the photo's area (sides rounded to 32,
    capped per variant), draw the red outline at that size with the matching sentence, and resize
    the result back to the photo's aspect ratio. Without a mask (or an empty one) it runs on the
    instruction alone. Returns (image, gpu seconds, variant, outlined)."""
    name = _resolve_variant(variant)
    instruction = (instruction or "").strip()
    area = photo.width * photo.height
    side = min(round(area**0.5), VARIANTS[name]["max_side"])
    work_size = _canvas(photo.width / photo.height, side)
    work = photo.resize(work_size, Image.LANCZOS)
    outlined = False
    if mask is not None:
        mask = mask.resize(work_size, Image.NEAREST)
        if (np.array(mask) > 127).any():
            work = draw_outline(work, mask)
            instruction = outline_prompt(instruction or DEFAULT_INSTRUCTION)
            outlined = True
    if not instruction:
        raise gr.Error("Paint over the object to remove, or describe it.")
    try:
        result, seconds = _render(instruction, name, seed, work, side)
    except (RuntimeError, ValueError) as error:
        raise gr.Error(str(error).split("\n")[0])
    scale = side / area**0.5
    out_size = (round(photo.width * scale), round(photo.height * scale)) if scale < 1 else photo.size
    return result.resize(out_size, Image.LANCZOS), seconds, name, outlined


# ------------------------------------------------------------------ helpers


def _seed(seed) -> int:
    try:
        return min(max(int(float(seed)), 0), MAX_SEED)
    except (TypeError, ValueError):
        return 42


def _resolve_seed(seed, randomize) -> int:
    return random.randint(0, MAX_SEED) if randomize else _seed(seed)


def _requested_path(value) -> str:
    """Turn whatever the canvas/API client sends into a path."""
    if isinstance(value, str):
        return value
    if not isinstance(value, dict):
        raise gr.Error(f"Expected an image, got {type(value).__name__}")
    if path := value.get("path"):
        return path
    url = value.get("url") or ""
    if "/gradio_api/file=" in url:
        return unquote(url.split("/gradio_api/file=", 1)[1])
    if url.startswith(("http://", "https://")):
        raise gr.Error("Upload the image instead of passing a URL.")
    return unquote(urlparse(url).path)


def _local_path(value) -> str:
    """A path this process may open: only files inside gradio's upload folder,
    where /gradio_api/upload puts uploads and _publish puts results."""
    path = os.path.realpath(_requested_path(value))
    folder = os.path.realpath(get_upload_folder())
    if os.path.commonpath([path, folder]) != folder:
        raise gr.Error("Upload the image first.")
    return path


def _empty_media(value) -> bool:
    """True when the canvas/API sent no file: None, '', or a dict with no path/url."""
    if value is None or value == "":
        return True
    if isinstance(value, Image.Image):
        return False
    if isinstance(value, str):
        return not value.strip()
    if isinstance(value, dict):
        return not (value.get("path") or value.get("url"))
    return True


def _as_instruction(value) -> str:
    if value is None or isinstance(value, bool):
        return ""
    if isinstance(value, str):
        return value
    if isinstance(value, dict):
        for key in ("text", "value", "prompt", "out", "in"):
            item = value.get(key)
            if isinstance(item, str) and item.strip():
                return item
        return ""
    return str(value)


def _open_rgb(value) -> Image.Image:
    if isinstance(value, Image.Image):
        return value.convert("RGB")
    return ImageOps.exif_transpose(Image.open(_local_path(value))).convert("RGB")


def _open_mask(value):
    """A mask as grayscale, white = the object. Transparent pixels count as black. None when absent."""
    if _empty_media(value):
        return None
    try:
        image = value if isinstance(value, Image.Image) else ImageOps.exif_transpose(Image.open(_local_path(value)))
    except (gr.Error, OSError, ValueError, FileNotFoundError):
        return None
    if "A" in image.getbands() or image.mode == "P":
        image = Image.alpha_composite(Image.new("RGBA", image.size, (0, 0, 0, 255)), image.convert("RGBA"))
    return image.convert("L")


def _as_data_uri(image: Image.Image) -> str:
    """Inline a result image so the studio page never needs a file route."""
    buffer = io.BytesIO()
    image.save(buffer, format="PNG")
    return "data:image/png;base64," + base64.b64encode(buffer.getvalue()).decode()


def _publish(image: Image.Image) -> str:
    """Save a result into gradio's upload folder, which it already serves, so the workflow canvas can render it."""
    folder = get_upload_folder()
    os.makedirs(folder, exist_ok=True)
    path = os.path.join(folder, f"pruna_{os.urandom(8).hex()}.png")
    image.save(path)
    return path


def _as_media(path: str) -> dict:
    """The dict shape the workflow canvas renders media from."""
    return {
        "path": path,
        "url": f"/gradio_api/file={client_utils.encode_file_path(path)}",
        "is_file": True,
    }


# ------------------------------------------------------------------ workflow fn


def pq_remove(image, prompt="", variant: str = DEFAULT_VARIANT, seed: str = "42", mask=None) -> dict:
    """Remove an object with a mask outline, a text instruction, or both (no CFG)."""
    if _empty_media(image):
        raise gr.Error("Connect a source image first.")
    result, _, _, _ = _remove(
        _open_rgb(image), _open_mask(mask), _as_instruction(prompt), variant, _seed(seed)
    )
    return _as_media(_publish(result))


# ------------------------------------------------------------------ studio server

server = gr.Server(title="P-Q-Image-2.1-Remove-Object")

EXAMPLE_FILES = {f"{eid}_{kind}" for eid, _, _ in EXAMPLES for kind in ("photo.jpg", "mask.png")}


@server.get("/", response_class=HTMLResponse)
def _studio_page():
    with open(os.path.join(APP_DIR, "studio.html"), encoding="utf-8") as handle:
        page = handle.read()
    config = {
        "examples": [{"id": eid, "label": label, "prompt": prompt} for eid, label, prompt in EXAMPLES],
        "variants": VARIANT_CHOICES,
        "defaultVariant": DEFAULT_VARIANT,
    }
    return page.replace("__STUDIO_CONFIG__", json.dumps(config).replace("</", "<\\/"))


@server.get("/pruna-logo.png")
def _logo():
    return FileResponse(os.path.join(APP_DIR, "pruna-logo.png"), media_type="image/png")


@server.get("/examples/{name}")
def _example_file(name: str):
    if name not in EXAMPLE_FILES:
        raise HTTPException(status_code=404)
    return FileResponse(os.path.join(EXAMPLES_DIR, name))


@server.api(name="remove", concurrency_limit=1)
def api_remove(
    source: str,
    mask: str = "",
    instruction: str = "",
    variant: str = DEFAULT_VARIANT,
    seed: float = 42,
    randomize: bool = True,
) -> dict:
    """Remove the masked object (white = remove) with P-Q-Image-2.1-Remove-Object."""
    if not source:
        raise gr.Error("Upload a photo first.")
    seed_value = _resolve_seed(seed, randomize)
    image, seconds, name, outlined = _remove(_open_rgb(source), _open_mask(mask), instruction, variant, seed_value)
    mode = "outline" if outlined else "prompt only"
    info = f"{image.width}×{image.height}, seed {seed_value} · {name} · {mode} · {seconds:.1f} s on the GPU"
    return {"image_data": _as_data_uri(image), "seed": seed_value, "info": info}


# The @server.api endpoints are deferred until Server.launch(), which also fires
# the ZeroGPU registration hook (pack weights + startup report) through the
# patched Blocks.launch path. Launch on a loopback port for registration, then
# serve the same app via uvicorn.
_port_env = os.environ.get("GRADIO_SERVER_PORT")
os.environ["GRADIO_SERVER_PORT"] = "7861"
try:
    server.launch(
        prevent_thread_lock=True,
        server_name="127.0.0.1",
        server_port=7861,
        quiet=True,
        show_error=False,
        ssr_mode=False,
    )
finally:
    if _port_env is None:
        del os.environ["GRADIO_SERVER_PORT"]
    else:
        os.environ["GRADIO_SERVER_PORT"] = _port_env

try:
    server.close()
except Exception:
    pass

# ------------------------------------------------------------------ workflow canvas
# Mounted only after the registration launch above. Mounting first starts the
# workflow's queue on the loopback server's event loop, and close() then strands
# it: jobs are accepted on port 7860 but never run.

workflow_demo = gr.Workflow(graph="workflow.json", bind={"pq_remove": pq_remove})
workflow_demo.title = "P-Q-Image-2.1-Remove-Object Workflow"

server = gr.mount_gradio_app(server, workflow_demo, path="/workflow", ssr_mode=False)

if __name__ == "__main__":
    uvicorn.run(server, host="0.0.0.0", port=7860)