Wire example click-to-fill with full generate() defaults; add in-process AoTI compilation of transformer blocks

app.py

@@ -17,6 +17,7 @@ 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"
 
@@ -26,6 +27,64 @@ pipe = JoyImageEditPlusPipeline.from_pretrained(
 )
 pipe.enable_model_cpu_offload()
 
+# A minimal, representative call used both to warm the pipeline and to capture
+# example inputs for ahead-of-time (AoTI) compilation of the repeated
+# transformer block. Kept tiny (1 step) so it's cheap during startup.
+_WARMUP_IMAGE = Image.new("RGB", (512, 512), (127, 127, 127))
+
+
+def _warmup_run():
+    """Run a single 1-step pipeline pass (used inside aoti_capture)."""
+    pipe(
+        images=[_WARMUP_IMAGE],
+        prompt="a photo",
+        negative_prompt="",
+        height=512,
+        width=512,
+        num_inference_steps=1,
+        guidance_scale=1.0,
+        generator=torch.Generator(device="cpu").manual_seed(0),
+    )
+
+
+@spaces.GPU(duration=1500)
+def compile_transformer_blocks():
+    """Ahead-of-time compile the repeated transformer block and apply it.
+
+    The transformer is a stack of identical ``JoyImageEditPlusTransformerBlock``
+    modules (``pipe.transformer.double_blocks``). We capture one block's real
+    inputs from a 1-step run, export + AoTI-compile that single block graph, and
+    swap the compiled forward into every block. Weights stay runtime inputs, so
+    one compiled graph serves all layers. This removes eager per-step overhead
+    and materially speeds up inference. torch.compile (JIT) is unsupported on
+    ZeroGPU, so this ahead-of-time path is the supported optimization.
+    """
+    blocks = pipe.transformer.double_blocks
+    block = blocks[0]
+
+    with spaces.aoti_capture(block) as call:
+        _warmup_run()
+
+    exported = torch.export.export(
+        block,
+        args=call.args,
+        kwargs=call.kwargs,
+        strict=False,
+    )
+    compiled = spaces.aoti_compile(exported)
+
+    for blk in blocks:
+        spaces.aoti_apply(compiled, blk)
+
+    return len(blocks)
+
+
+try:
+    _n = compile_transformer_blocks()
+    print(f"AoTI: compiled and applied to {_n} transformer blocks")
+except Exception as e:  # noqa: BLE001 -- keep the Space running if AoTI fails
+    print(f"AoTI compile failed ({e!r}); running eager")
+
 
 @spaces.GPU(duration=300)
 def generate(
@@ -177,11 +236,30 @@ with gr.Blocks() as demo:
     )
 
     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"],
+            [
+                ["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,
         ],
-        inputs=[input_gallery, prompt],
         outputs=[output_image, seed],
         fn=generate,
         cache_examples=True,