Upload app.py with huggingface_hub

app.py

@@ -27,63 +27,34 @@ 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))
+# AoTI: load the pre-compiled repeated transformer block graph from the Hub and
+# apply it to every ``JoyImageEditPlusTransformerBlock``
+# (``pipe.transformer.double_blocks``, 40 of them). The graph is compiled
+# 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``.
+# Weights stay runtime inputs, so the single compiled graph serves all blocks —
+# and no torch.compile happens in this live demo process. 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 snapshot_download
 
-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),
+    _aoti_dir = snapshot_download(
+        repo_id=AOTI_DATASET_REPO,
+        repo_type="dataset",
+        allow_patterns="package/*",
     )
-
-
-@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,
+    spaces.aoti_load_from_package_dir(
+        pipe.transformer.double_blocks,
+        f"{_aoti_dir}/package",
     )
-    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")
+    print(
+        f"AoTI: loaded compiled block graph from {AOTI_DATASET_REPO} and applied "
+        f"to {len(pipe.transformer.double_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)