Spaces:
Running on Zero
Running on Zero
| """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"] | |
| 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")} | |
| 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("</", "<\\/")) | |
| def _logo(): | |
| return FileResponse(os.path.join(APP_DIR, "pruna-logo.png"), media_type="image/png") | |
| def _example_file(name: str): | |
| if name not in EXAMPLE_FILES: | |
| raise HTTPException(status_code=404) | |
| return FileResponse(os.path.join(EXAMPLES_DIR, name)) | |
| 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) | |