"""JoyAI-Image Edit Plus — multi-image instruction-guided editing demo. Loads the jdopensource/JoyAI-Image-Edit-Plus-Diffusers model and exposes a Gradio interface where visitors provide one or more reference images and a text instruction; the model generates a new image that combines elements from the references according to the instruction. """ import os os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") import spaces # noqa: E402 -- must precede torch import import random # noqa: E402 import gradio as gr # noqa: E402 import torch # noqa: E402 from diffusers import JoyImageEditPlusPipeline # noqa: E402 from PIL import Image # noqa: E402 MODEL_ID = "jdopensource/JoyAI-Image-Edit-Plus-Diffusers" pipe = JoyImageEditPlusPipeline.from_pretrained( MODEL_ID, torch_dtype=torch.bfloat16, ) pipe.to("cuda") # AoTI: instead of compiling the repeated transformer block in this live demo # process, download the pre-compiled graph produced offline by the one-shot # Space multimodalart/joyai-image-edit-plus-aoti-export and published to the # public dataset below as ``package/root/package.pt2``. We load that compiled # graph and apply it to every ``JoyImageEditPlusTransformerBlock`` # (``pipe.transformer.double_blocks``, 40 of them). Weights stay runtime inputs, # so the single compiled graph serves all blocks and no torch.compile runs here. # # Application happens at module scope (in the main process) so every forked # ZeroGPU worker inherits the patched blocks. Each block is wrapped in its own # ``ZeroGPUCompiledModel`` keyed by that block's weights (``ZeroGPUWeights`` # moves the constants onto CUDA when the worker is forked, so cpu-offload's # CPU-resident weights are handled correctly). Falls back to eager execution if # the artifact can't be downloaded/applied. AOTI_DATASET_REPO = "multimodalart/joyai-image-edit-plus-aoti-pt2" try: from huggingface_hub import hf_hub_download from spaces.zero.torch.aoti import ZeroGPUCompiledModel, ZeroGPUWeights _pt2_path = hf_hub_download( repo_id=AOTI_DATASET_REPO, repo_type="dataset", filename="package/root/package.pt2", ) _blocks = pipe.transformer.double_blocks # Build a single compiled model from the first block's weights and swap it # into every block — the same application the in-process compile did (which # applied one ``ZeroGPUCompiledModel`` to all blocks). to_cuda=True routes the # weights through ZeroGPU's fake-cuda pack machinery at module scope so the # constants are streamed onto real CUDA in the forked GPU worker (matching # what the compiled graph expects, even though the pipeline uses cpu-offload). _weights = ZeroGPUWeights(_blocks[0].state_dict(), to_cuda=True) _compiled = ZeroGPUCompiledModel(_pt2_path, _weights) for _blk in _blocks: spaces.aoti_apply(_compiled, _blk) print( f"AoTI: loaded compiled block graph from {AOTI_DATASET_REPO} and applied " f"to {len(_blocks)} transformer blocks" ) except Exception as e: # noqa: BLE001 -- keep the Space running if AoTI load fails print(f"AoTI load failed ({e!r}); running eager") @spaces.GPU(duration=300) def generate( images, prompt: str, negative_prompt: str, num_inference_steps: int, guidance_scale: float, seed: int, randomize_seed: bool, progress=gr.Progress(track_tqdm=True), ): """Edit/generate an image from multiple reference images and a text instruction. Args: images: One or more reference images (1-6 supported). prompt: Text instruction describing the desired edit or composition. negative_prompt: Negative prompt to guide what to avoid. num_inference_steps: Number of denoising steps (30 recommended). guidance_scale: Classifier-free guidance scale (4.0 recommended). seed: RNG seed for reproducibility. randomize_seed: If True, pick a random seed each run. """ if images is None or (isinstance(images, list) and len(images) == 0): raise gr.Error("Please provide at least one reference image.") if not prompt or not prompt.strip(): raise gr.Error("Please provide a text instruction.") # Gradio Gallery returns list of (filepath, caption) tuples, dicts, or strings from PIL import Image if isinstance(images, list): pil_images = [] for img in images: if isinstance(img, dict): path = img.get("image") or img.get("path") if path: pil_images.append(Image.open(path).convert("RGB")) elif isinstance(img, (tuple, list)) and len(img) > 0: path = img[0] pil_images.append(Image.open(path).convert("RGB")) elif isinstance(img, str): pil_images.append(Image.open(img).convert("RGB")) elif isinstance(img, Image.Image): pil_images.append(img.convert("RGB")) elif isinstance(images, str): pil_images = [Image.open(images).convert("RGB")] else: pil_images = [images.convert("RGB")] if len(pil_images) > 6: pil_images = pil_images[:6] if randomize_seed: seed = random.randint(0, 2**31 - 1) # Determine output resolution from the last reference image target_h, target_w = pipe.vae_image_processor.get_default_height_width(pil_images[-1]) generator = torch.Generator(device="cpu").manual_seed(int(seed)) result = pipe( images=pil_images, prompt=prompt, negative_prompt=negative_prompt, height=target_h, width=target_w, num_inference_steps=int(num_inference_steps), guidance_scale=float(guidance_scale), generator=generator, ) output_image = result.images[0] return output_image, seed CSS = """ #col-container { max-width: 1100px; margin: 0 auto; } .dark .gradio-container { color: var(--body-text-color); } """ with gr.Blocks() as demo: gr.Markdown( """ # JoyAI-Image Edit Plus Multi-image instruction-guided editing — provide reference images and a text instruction to generate a new image combining elements from the references. [Model card](https://huggingface.co/jdopensource/JoyAI-Image-Edit-Plus-Diffusers) """ ) with gr.Row(elem_id="col-container"): with gr.Column(scale=1): input_gallery = gr.Gallery( label="Reference images", show_label=True, columns=3, height=240, object_fit="contain", file_types=["image"], type="filepath", ) prompt = gr.Textbox( label="Edit instruction", placeholder="e.g. The woman is lovingly holding the cute puppy in her arms", lines=2, ) run_btn = gr.Button("Generate", variant="primary") with gr.Accordion("Advanced settings", open=False): negative_prompt = gr.Textbox( label="Negative prompt", value="low quality, blurry, deformed", lines=2, ) num_inference_steps = gr.Slider( label="Inference steps", minimum=1, maximum=100, step=1, value=30, ) guidance_scale = gr.Slider( label="Guidance scale", minimum=1.0, maximum=20.0, step=0.1, value=4.0, ) randomize_seed = gr.Checkbox(label="Randomize seed", value=True) seed = gr.Number(label="Seed", value=42, precision=0) with gr.Column(scale=1): output_image = gr.Image(label="Result", show_label=True, height=420) run_btn.click( fn=generate, inputs=[ input_gallery, prompt, negative_prompt, num_inference_steps, guidance_scale, seed, randomize_seed, ], outputs=[output_image, seed], api_name="generate", ) gr.Examples( # Each row supplies a full set of default values for every generate() # input, in the same order as the `inputs` list below, so that clicking # an example both populates the controls and calls generate with those # defaults. examples=[ [ ["examples/input_0.png", "examples/input_1.png"], "The woman is lovingly holding the cute puppy in her arms", "low quality, blurry, deformed", 30, 4.0, 42, False, ], ], inputs=[ input_gallery, prompt, negative_prompt, num_inference_steps, guidance_scale, seed, randomize_seed, ], outputs=[output_image, seed], fn=generate, cache_examples=True, cache_mode="lazy", ) demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS)