Spaces:
Running on Zero
Running on Zero
| """P-H3-Image on ZeroGPU (`xlarge`: a full RTX PRO 6000 Blackwell, 96 GB). | |
| The pipeline is loaded once at startup, one component at a time (the host has ~97 GB), and slimmed to fit the slice | |
| with p_h3_image_fast.py (kernels only, no torch.compile): AdaLN tables instead of the transformer's AdaLN projections, | |
| FP8 linears everywhere, fused Triton DiT blocks, batched VAE tiles, and for the prompt upsampler a weight-only FP8 GEMV, | |
| decode attention and speculative decoding with its MTP head in one CUDA graph per step. | |
| """ | |
| import json | |
| import os | |
| import random | |
| import tempfile | |
| import time | |
| os.environ.setdefault("DIFFUSERS_TRUST_REMOTE_KERNELS", "true") | |
| os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") | |
| import gradio as gr | |
| import spaces | |
| import torch | |
| from huggingface_hub import snapshot_download | |
| from PIL import Image, ImageOps | |
| from torchvision.io import encode_jpeg | |
| import p_h3_image_fast as fast | |
| from diffusers import ModularPipeline | |
| from examples import EXAMPLES | |
| from theme import CSS, HEAD, THEME, render_footer, render_header | |
| REPO, REVISION = "PrunaAI/P-H3-Image", "ee4a685de1fbd4337a50bcf9b3c407e4f8463d49" | |
| # Canvases (width, height), multiples of 32: 4 and 8 MP as tested, 16 MP from the same ratios. | |
| RATIOS = {"1:1": 1.0, "4:3": 4 / 3, "3:4": 3 / 4, "3:2": 3 / 2, "2:3": 2 / 3, "16:9": 16 / 9, "9:16": 9 / 16} | |
| T2I_RESOLUTIONS, EDIT_RESOLUTIONS = ["4 MP", "8 MP", "16 MP"], ["4 MP", "8 MP"] | |
| def _canvas(ratio: float, megapixels: int) -> tuple[int, int]: | |
| pixels = megapixels * 1024**2 | |
| return round((pixels * ratio) ** 0.5 / 32) * 32, round((pixels / ratio) ** 0.5 / 32) * 32 | |
| SIZES = { | |
| "4 MP": {"1:1": (2048, 2048), "4:3": (2368, 1760), "3:4": (1760, 2368), "3:2": (2496, 1664), "2:3": (1664, 2496), | |
| "16:9": (2720, 1536), "9:16": (1536, 2720)}, | |
| "8 MP": {"1:1": (2912, 2912), "4:3": (3360, 2496), "3:4": (2496, 3360), "3:2": (3552, 2368), "2:3": (2368, 3552), | |
| "16:9": (3872, 2176), "9:16": (2176, 3872)}, | |
| "16 MP": {name: _canvas(ratio, 16) for name, ratio in RATIOS.items()}, | |
| } | |
| MAX_IMAGES = 8 | |
| # All references together stay within 8 MP: one image keeps up to 8 MP, eight get 1 MP each. | |
| REFERENCE_BUDGET = 8 * 1024**2 | |
| def load_pipeline(): | |
| path = snapshot_download(REPO, revision=REVISION, max_workers=16) | |
| with open(os.path.join(os.path.dirname(__file__), "adaln_calls.json")) as handle: | |
| calls = json.load(handle) | |
| pipe, mtp = fast.load_slim_pipeline(path, calls, log=lambda message: print(message, flush=True)) | |
| # The same components without the prompt upsampler step, for requests that skip upsampling. | |
| blocks = ModularPipeline.from_pretrained(path).blocks | |
| blocks.sub_blocks.pop("prompt_upsampler") | |
| plain = blocks.init_pipeline(path) | |
| plain.update_components(**{name: getattr(pipe, name) for name in plain.component_names | |
| if getattr(pipe, name, None) is not None}) | |
| return pipe, plain, mtp | |
| PIPE, PLAIN, MTP = load_pipeline() | |
| UPSAMPLER_GENERATE = PIPE.prompt_upsampler.generate | |
| STATE = {"kernels": False, "decoder": None, "owner": None} | |
| def _prepare_gpu(): | |
| """Per GPU call. Once per worker process: the upsampler's hub kernels, and the speculative decoder with its CUDA | |
| graph (captured on the first request, then replayed; rebuilt if the process or device changes).""" | |
| upsampler = PIPE.prompt_upsampler | |
| if not STATE["kernels"]: | |
| from kernels import Mode | |
| upsampler.set_use_kernels(True, mode=Mode.INFERENCE) | |
| STATE["kernels"] = True | |
| # The graph bakes in weight addresses: rebuild if ZeroGPU placed the weights elsewhere since the capture. | |
| weights = (upsampler.lm_head.weight, MTP.fc.weight, upsampler.model.language_model.embed_tokens.weight, | |
| upsampler.model.language_model.layers[0].mlp.down_proj.weight) | |
| owner = (os.getpid(), torch.cuda.current_device(), tuple(w.data_ptr() for w in weights)) | |
| if STATE["decoder"] is None or STATE["owner"] != owner: | |
| upsampler.generate = UPSAMPLER_GENERATE | |
| STATE["decoder"], STATE["owner"] = fast.MTPGreedyGenerate(upsampler, MTP), owner | |
| decoder = STATE["decoder"] | |
| timing = {} | |
| def generate(*args, **kwargs): | |
| torch.cuda.synchronize() | |
| begin = time.perf_counter() | |
| out = decoder.generate(*args, **kwargs) | |
| torch.cuda.synchronize() | |
| timing["upsampler"] = time.perf_counter() - begin | |
| timing["tokens"] = decoder.stats.get("tokens", 0) | |
| return out | |
| upsampler.generate = generate | |
| return timing | |
| def _duration(mode, prompt, images, aspect, resolution, upsample, seed): | |
| # Measured on a full card: ~6 / 12 / 25 s for 4 / 8 / 16 MP text-to-image, edits a bit more per image. Shorter | |
| # requests queue faster on ZeroGPU. | |
| base = {"4 MP": 25, "8 MP": 35, "16 MP": 60}[resolution] | |
| return base + (10 + 5 * len(images) if mode == "edit" else 0) | |
| def _render(mode, prompt, images, aspect, resolution, upsample, seed): | |
| timing = _prepare_gpu() | |
| torch.cuda.synchronize() | |
| start = time.perf_counter() | |
| kwargs = {"prompt": prompt, "generator": torch.Generator().manual_seed(seed)} | |
| if mode == "edit": | |
| # The first image's aspect ratio at the chosen size. | |
| width, height = images[0].size | |
| kwargs["image"] = images | |
| kwargs["width"], kwargs["height"] = _canvas(width / height, int(resolution.split()[0])) | |
| else: | |
| kwargs["width"], kwargs["height"] = SIZES[resolution][aspect] | |
| result = (PIPE if upsample else PLAIN)(**kwargs, output_type="pt", output=["images", "prompt"]) | |
| path, size = _to_jpeg(result["images"]) | |
| total = time.perf_counter() - start | |
| return path, size, result["prompt"], {"total": total, **timing} | |
| def _to_jpeg(images: torch.Tensor) -> tuple[str, tuple[int, int]]: | |
| """The decoded image as a JPEG file: uint8 and nvjpeg on the GPU, so only the encoded bytes leave the worker.""" | |
| pixels = (images[0].clamp(0, 1) * 255).round().to(torch.uint8).contiguous() | |
| try: | |
| data = encode_jpeg(pixels, quality=93) | |
| except Exception: # noqa: BLE001 - no nvjpeg: encode on the CPU | |
| data = encode_jpeg(pixels.cpu(), quality=93) | |
| with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as handle: | |
| handle.write(data.cpu().numpy().tobytes()) | |
| return handle.name, (pixels.shape[2], pixels.shape[1]) | |
| def _warmup(): | |
| """Triton compiles and autotunes, and the hub kernels load, before the first visitor.""" | |
| image = Image.new("RGB", (1024, 1024), (120, 140, 160)) | |
| timing = _prepare_gpu() | |
| PIPE(prompt="a lighthouse on a cliff at dusk", width=1024, height=1024, generator=torch.Generator().manual_seed(0), | |
| output="images") | |
| PIPE(prompt="make it night", image=[image], width=1024, height=1024, generator=torch.Generator().manual_seed(0), | |
| output="images") | |
| return timing | |
| try: | |
| _warmup() | |
| except Exception as error: # noqa: BLE001 - the demo still works, the first request just pays the warmup | |
| print(f"warmup failed: {error!r}", flush=True) | |
| def _info(timing, size, seed): | |
| upsample, tokens = timing.get("upsampler", 0.0), timing.get("tokens", 0) | |
| rate = f", {tokens} tokens at {tokens / upsample:.0f} tokens/s" if upsample and tokens else "" | |
| head = f"{size[0]}×{size[1]}, seed {seed} · {timing['total']:.1f} s on the GPU" | |
| if not upsample: | |
| return f"{head} (no prompt upsampling)" | |
| return f"{head} ({upsample:.1f} s prompt upsampling{rate}, {timing['total'] - upsample:.1f} s rendering)" | |
| def text_to_image(prompt, aspect, resolution, upsample): | |
| if not prompt or not prompt.strip(): | |
| raise gr.Error("Write a prompt first.") | |
| seed = random.randint(0, 2**31 - 1) | |
| try: | |
| path, size, upsampled, timing = _render("t2i", prompt.strip(), [], aspect, resolution, upsample, seed) | |
| except RuntimeError as error: | |
| raise gr.Error(str(error).split("\n")[0]) | |
| return path, upsampled, _info(timing, size, seed) | |
| def _fit(image: Image.Image, pixels: float) -> Image.Image: | |
| """Downscale to at most `pixels` (the pipeline keeps references at their own size, from 1 MP up).""" | |
| scale = (pixels / (image.width * image.height)) ** 0.5 | |
| if scale >= 1: | |
| return image | |
| return image.resize((max(1, round(image.width * scale)), max(1, round(image.height * scale))), Image.LANCZOS) | |
| def edit(gallery, prompt, resolution, upsample): | |
| images = [] | |
| for item in gallery or []: | |
| item = item[0] if isinstance(item, (tuple, list)) else item | |
| item = Image.open(item) if isinstance(item, str) else item | |
| images.append(ImageOps.exif_transpose(item).convert("RGB")) | |
| if not images: | |
| raise gr.Error("Add at least one image; the first one is edited.") | |
| if len(images) > MAX_IMAGES: | |
| raise gr.Error(f"At most {MAX_IMAGES} images.") | |
| if not prompt or not prompt.strip(): | |
| raise gr.Error("Describe the edit first.") | |
| images = [_fit(image, REFERENCE_BUDGET / len(images)) for image in images] | |
| seed = random.randint(0, 2**31 - 1) | |
| try: | |
| path, size, upsampled, timing = _render("edit", prompt.strip(), images, None, resolution, upsample, seed) | |
| except (RuntimeError, ValueError) as error: | |
| raise gr.Error(str(error).split("\n")[0]) | |
| return path, upsampled, _info(timing, size, seed) | |
| with gr.Blocks(title="P-H3-Image", fill_width=True) as demo: | |
| render_header() | |
| with gr.Tabs(elem_classes="main-tabs"): | |
| with gr.Tab("Text to image"): | |
| with gr.Row(): | |
| with gr.Column(): | |
| t_prompt = gr.Textbox(label="Prompt", lines=4) | |
| with gr.Row(): | |
| t_aspect = gr.Dropdown(list(RATIOS), value="4:3", label="Aspect ratio", filterable=False) | |
| t_resolution = gr.Radio(T2I_RESOLUTIONS, value="8 MP", label="Resolution") | |
| t_upsample = gr.Checkbox(value=True, label="Prompt upsampling") | |
| t_button = gr.Button("Generate", variant="primary") | |
| with gr.Column(): | |
| t_image = gr.Image(label="Result", type="filepath") | |
| t_info = gr.Markdown() | |
| with gr.Accordion("Upsampled prompt", open=False): | |
| t_upsampled = gr.Textbox(show_label=False, lines=10) | |
| gr.Examples([[prompt, aspect] for _, aspect, prompt in EXAMPLES], inputs=[t_prompt, t_aspect], | |
| example_labels=[label for label, _, _ in EXAMPLES]) | |
| t_button.click(text_to_image, [t_prompt, t_aspect, t_resolution, t_upsample], [t_image, t_upsampled, t_info]) | |
| with gr.Tab("Edit"): | |
| with gr.Row(): | |
| with gr.Column(): | |
| e_images = gr.Gallery(label=f"Images (the first one is edited, up to {MAX_IMAGES})", type="pil", | |
| columns=4, height=260, interactive=True) | |
| e_prompt = gr.Textbox(label="Edit", lines=3, placeholder="e.g. put the fox from <Picture 2> on the bench") | |
| e_resolution = gr.Radio(EDIT_RESOLUTIONS, value="8 MP", label="Resolution", | |
| info="The first image's aspect ratio at this size.") | |
| e_upsample = gr.Checkbox(value=True, label="Prompt upsampling") | |
| e_button = gr.Button("Edit", variant="primary") | |
| with gr.Column(): | |
| e_image = gr.Image(label="Result", type="filepath") | |
| e_info = gr.Markdown() | |
| with gr.Accordion("Upsampled prompt", open=False): | |
| e_upsampled = gr.Textbox(show_label=False, lines=10) | |
| e_button.click(edit, [e_images, e_prompt, e_resolution, e_upsample], [e_image, e_upsampled, e_info]) | |
| render_footer() | |
| demo.queue(max_size=20).launch(theme=THEME, css=CSS, head=HEAD) | |