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

The Space currently returns pure noise for every generation (verified live: 30 steps, guidance 4.0, seed 42 on the model card's own example returns an 832x1152 image of RGB noise).

### Root cause: the AoTI graph runs all 40 blocks with block 0's weights

`package.pt2` is exported with the block's weights as *runtime* constants, so one compiled graph can serve all 40 `JoyImageEditPlusTransformerBlock`s — but each block still has to be bound to **its own** weights. The current code builds a single `ZeroGPUCompiledModel` from `_blocks[0].state_dict()` and applies that same instance to every block:

```python
_weights = ZeroGPUWeights(_blocks[0].state_dict(), to_cuda=True)
_compiled = ZeroGPUCompiledModel(_pt2_path, _weights)
for _blk in _blocks:
spaces.aoti_apply(_compiled, _blk)
```

`ZeroGPUCompiledModel.__call__` loads `self.weights.constants_map` once and reuses it, so the transformer effectively becomes block 0 repeated 40 times — no error, just noise.

The supported API does the right thing: given a `ModuleList`, `spaces.aoti_load_from_package_dir` patches every block with one shared `LazyAOTIModel` while keeping that block's own `state_dict` as the constants to load. (Same mechanism as `spaces.aoti_blocks_load`; used directly here because the artifact lives in a dataset repo.)

### Fallout: the demo was never actually resident on the GPU

`spaces.aoti_apply` calls `drain_module_parameters`, which replaced all 40 blocks' parameters with empty tensors — silently dropping ~29GB from the ZeroGPU pack. That is the only reason the Space fit the default 48GB MIG slice. With the weights correctly retained the pack is 50.3GB, which does not fit, so the GPU function now requests `size="xlarge"` (full 96GB card). Verified: without it the first CUDA op OOMs (`MIG 2g.48gb`, 47.4GiB).

### Two robustness fixes found along the way

- **`requirements.txt` pins `diffusers` to a moving branch.** The revision the Space is currently running is not the revision a rebuild would install. On the branch tip `JoyImageEditPlusPipeline` renamed `vae_image_processor` -> `image_processor` and changed the block's `image_rotary_emb` argument from `((cos, sin), None)` to `(cos, sin)` — the first breaks `app.py`, the second invalidates `package.pt2`. Pinned to `3fa516a2`, the branch tip the Space and the AoTI artifact were both built from.
- **`app.py` no longer calls `pipe.vae_image_processor`.** `height`/`width` are left unset; the pipeline derives the same 1024-base bucket from the last reference image internally, so this works on either pipeline revision.

### Verification

Tested end-to-end on a duplicate Space with the same inputs, prompt and seed as the `examples/` row. Before: RGB noise. After: a coherent generation matching the reference output.

*Alternative if `xlarge` (2x visitor quota) is not wanted:* drop the AoTI block and go back to `pipe.enable_model_cpu_offload()`. Peak VRAM then stays under the 48GB slice, at the cost of the compiled graph and per-call weight transfers. Happy to switch this PR to that shape instead.

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