[Admin] Fix garbled output: bind each transformer block to its own AoTI weights

app.py

@@ -30,40 +30,39 @@ 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.
+# public dataset below as ``package/root/package.pt2``.
+#
+# The graph was exported with the block's weights left as *runtime* constants,
+# so the single compiled graph serves all 40
+# ``JoyImageEditPlusTransformerBlock``s (``pipe.transformer.double_blocks``) --
+# but each block must be bound to its OWN weights. ``aoti_load_from_package_dir``
+# does exactly that: given a ``ModuleList`` it patches every block with the one
+# shared ``LazyAOTIModel`` while keeping that block's own state_dict as the
+# constants to load.
 #
 # 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
+# ZeroGPU worker inherits the patched blocks. 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
+    from pathlib import Path  # noqa: E402
+
+    from huggingface_hub import snapshot_download
 
-    _pt2_path = hf_hub_download(
-        repo_id=AOTI_DATASET_REPO,
-        repo_type="dataset",
-        filename="package/root/package.pt2",
+    _package_dir = (
+        Path(
+            snapshot_download(
+                repo_id=AOTI_DATASET_REPO,
+                repo_type="dataset",
+                allow_patterns="package/*",
+            )
+        )
+        / "package"
     )
 
     _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)
+    spaces.aoti_load_from_package_dir(_blocks, _package_dir)
 
     print(
         f"AoTI: loaded compiled block graph from {AOTI_DATASET_REPO} and applied "
@@ -73,7 +72,10 @@ except Exception as e:  # noqa: BLE001 -- keep the Space running if AoTI load fa
     print(f"AoTI load failed ({e!r}); running eager")
 
 
-@spaces.GPU(duration=300)
+# ``size="xlarge"`` (full 96GB card) is required, not optional: the pipeline is
+# ~50GB of bf16 weights (16B transformer + Qwen3-VL-8B text encoder), which
+# does not fit the default 48GB MIG slice.
+@spaces.GPU(duration=300, size="xlarge")
 def generate(
     images,
     prompt: str,
@@ -128,17 +130,15 @@ def generate(
     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))
 
+    # `height`/`width` are left unset: the pipeline picks the 1024-base bucket
+    # from the last reference image itself, which is what the model card
+    # recommends and avoids depending on the image-processor attribute name.
     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,

requirements.txt

@@ -1,4 +1,8 @@
-git+https://github.com/tangyanf/diffusers.git@add-joyimage-edit-plus
+# Pinned to the exact branch commit this Space was built against: it is the
+# revision the AoTI package.pt2 artifact was exported from, and later commits
+# on the moving branch changed both the block's call signature and the
+# pipeline's image-processor attribute name.
+git+https://github.com/tangyanf/diffusers.git@3fa516a2201609cf0488b992dca96b148d512240
 transformers>=4.57.0
 accelerate
 sentencepiece