app.py

"""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)