[Model] Add JoyAI-Image-Edit pipeline support (#4112)

Signed-off-by: nagisa-kun <1434936049@qq.com>
Signed-off-by: nagisa <1434936049@qq.com>
Signed-off-by: hurukawa <61525444+nagisa-kunhah@users.noreply.github.com>

.buildkite/common/ci_source_file_dependencies.yml

@@ -115,6 +115,10 @@ x-qwen-image: &qwen_image
   - vllm_omni/diffusion/models/dmd2/
   - vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py
 
+x-joy-image: &joy_image
+  - vllm_omni/diffusion/models/joy_image/
+  - vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py
+
 x-hunyuan-image3: &hunyuan_image3
   - vllm_omni/diffusion/models/hunyuan_image3/
   - vllm_omni/model_executor/models/hunyuan_image3/
@@ -387,6 +391,12 @@ source_file_dependencies:
     - tests/e2e/accuracy/test_qwen_image.py
     - tests/e2e/accuracy/test_diffusers_backend_similarity.py
 
+  diffusion_joy_image_accuracy:
+    - *joy_image
+    - tests/diffusion/models/joy_image/
+    - tests/e2e/accuracy/test_joy_image_edit.py
+    - recipes/JD/JoyAI-Image-Edit.md
+
   diffusion_qwen_image_perf:
     - *qwen_image
     - tests/dfx/perf/scripts/run_diffusion_benchmark.py

.buildkite/cuda/test-nightly.yml

@@ -268,6 +268,12 @@ steps:
           - pytest -s -v tests/e2e/accuracy/test_qwen_image.py -m "H100 and B200 and cards_1" --run-level full_model
           - pytest -s -v tests/e2e/accuracy/test_diffusers_backend_similarity.py -k '2i_matches_diffusers' -m "full_model and H100 and B200 and cards_1" --run-level full_model
 
+      - label: ":full_moon: JoyAI-Image-Edit · Accuracy Test"
+        source_file_dependencies: diffusion_joy_image_accuracy
+        timeout_in_minutes: 180
+        commands:
+          - pytest -s -v tests/e2e/accuracy/test_joy_image_edit.py -m "full_model and diffusion and H100 and B200 and cards_1" --run-level full_model
+
       - label: ":full_moon: HunyuanImage3-DIT · Accuracy Test"
         source_file_dependencies: diffusion_hunyuan_image3_accuracy
         timeout_in_minutes: 180

docs/models/supported_models.md

@@ -29,6 +29,7 @@ th {
 | `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` | ✅︎ | | | ✅︎ | [Published](https://recipes.vllm.ai/Qwen/Qwen-Image) |
 | `QwenImagePipeline` | Qwen-Image-2512 | `Qwen/Qwen-Image-2512` | ✅︎ | | | ✅︎ | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Qwen/Qwen-Image-2512.md) |
 | `QwenImageEditPipeline` | Qwen-Image-Edit | `Qwen/Qwen-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
+| `JoyImageEditPipeline` | JoyAI-Image-Edit | `jdopensource/JoyAI-Image-Edit-Diffusers` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/JD/JoyAI-Image-Edit.md) |
 | `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2509 | `Qwen/Qwen-Image-Edit-2509` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
 | `QwenImageLayeredPipeline` | Qwen-Image-Layered | `Qwen/Qwen-Image-Layered` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
 | `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2511 | `Qwen/Qwen-Image-Edit-2511` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |

docs/user_guide/diffusion/cpu_offload.md

@@ -42,6 +42,7 @@ higher-priority strategy is selected.
 | --- | --- | --- | :---: | :---: | :---: | --- |
 | Flux2Pipeline | `black-forest-labs/FLUX.2-dev` | `Flux2Transformer2DModel` | yes | yes | — | `transformer_blocks`, `single_transformer_blocks` |
 | LongCatImagePipeline | `meituan-longcat/LongCat-Image` | `LongCatImageTransformer2DModel` | — | yes | — | `transformer_blocks`, `single_transformer_blocks` |
+| JoyImageEditPipeline | `jdopensource/JoyAI-Image-Edit-Diffusers` | `JoyImageEditTransformer3DModel` | yes | yes | — | `double_blocks` |
 | NextStep11Pipeline | `stepfun-ai/NextStep-1.1` | `NextStepModel` | — | yes | — | `layers` |
 | OvisImagePipeline | `AIDC-AI/Ovis-Image-7B` | `OvisImageTransformer2DModel` | — | yes | — | `transformer` |
 | QwenImagePipeline | `Qwen/Qwen-Image` | `QwenImageTransformer2DModel` | yes | yes | — | `transformer_blocks` |

docs/user_guide/diffusion_features.md

@@ -152,6 +152,7 @@ The following tables show which models support each feature:
 | **HiDream-O1-Image**     |     ❌     |     ✅      |           ❌           |       ❌        |         ✅         |          ❌          |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
 | **HunyuanImage3**        |     ❌     |     ✅      |           ❌           |       ❌        |         ✅         |          ❌          |   ❌    |             ❌             |          ❌           |       ✅        |        ✅*        |
 | **Krea 2**               |     ❌     |     ✅      |           ❌           |       ❌        |         ❌         |          ❌          |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
+| **JoyAI-Image-Edit**    |     ❌     |     ❌      |           ❌           |       ✅        |         ❌         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **LongCat-Image**        |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **LongCat-Image-Edit**   |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **MammothModa2(T2I)**    |     ✅     |     ✅      |           ❌           |       ❌        |         ❌         |          ❌          |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |

docs/user_guide/examples/offline_inference/image_to_image.md

@@ -3,7 +3,7 @@
 Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/image_to_image>.
 
 
-This example edits an input image with `Qwen/Qwen-Image-Edit` using the `image_edit.py` CLI.
+This example edits an input image with `Qwen/Qwen-Image-Edit` and `JoyAI-Image-Edit` using the `image_edit.py` CLI.
 
 ## Local CLI Usage
 
@@ -26,6 +26,27 @@ python image_edit.py \
   --cfg-scale 4.0
 ```
 
+### JoyAI-Image-Edit
+
+JoyAI-Image-Edit uses the same `image_edit.py` CLI and supports one input
+image per request. Requested dimensions are snapped to the nearest supported
+Joy bucket; for square outputs, use `1024x1024`.
+
+```bash
+python image_edit.py \
+  --model jdopensource/JoyAI-Image-Edit-Diffusers \
+  --image input.png \
+  --prompt "Change the background to a clean studio while preserving the subject." \
+  --height 1024 \
+  --width 1024 \
+  --num-inference-steps 50 \
+  --guidance-scale 4.0 \
+  --output output_joyai_edit.png
+```
+
+JoyAI defaults to guidance scale `4.0`; pass `--guidance-scale` explicitly when
+you need a different value.
+
 ### Multiple Image Editing (Qwen-Image-Edit-2509)
 
 For multiple image inputs, use `Qwen/Qwen-Image-Edit-2509` or  `Qwen/Qwen-Image-Edit-2511`:
@@ -43,7 +64,7 @@ python image_edit.py \
 
 Key arguments:
 
-- `--model`: model name or path. Use `Qwen/Qwen-Image-Edit-2509` or later for multiple image support.
+- `--model`: model name or path. Use `Qwen/Qwen-Image-Edit-2509` or later for multiple image support. `jdopensource/JoyAI-Image-Edit-Diffusers` supports one input image.
 - `--image`: path(s) to the source image(s) (PNG/JPG, converted to RGB). Can specify multiple images.
 - `--prompt` / `--negative-prompt`: text description (string).
 - `--cfg-scale`: true classifier-free guidance scale (default: 4.0). Classifier-free guidance is enabled by setting cfg_scale > 1 and providing a negative_prompt. Higher guidance scale encourages images closely linked to the text prompt, usually at the expense of lower image quality.

docs/user_guide/examples/online_serving/image_to_image.md

@@ -3,10 +3,14 @@
 Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/image_to_image>.
 
 
-This example demonstrates how to deploy Qwen-Image-Edit model for online image editing service using vLLM-Omni.
+This example demonstrates how to deploy image-to-image models for online image editing service using vLLM-Omni.
+
+Supported models include Qwen-Image-Edit, BAGEL, JoyAI-Image-Edit, and other image-to-image pipelines.
 
 For **multi-image** input editing, use **Qwen-Image-Edit-2509** (QwenImageEditPlusPipeline) and send multiple images in the user message content.
 
+For **JoyAI-Image-Edit**, use `jdopensource/JoyAI-Image-Edit-Diffusers` and send exactly one input image.
+
 ## Start Server
 
 ### Basic Start
@@ -24,6 +28,12 @@ vllm serve Qwen/Qwen-Image-Edit --omni --port 8092
 vllm serve Qwen/Qwen-Image-Edit-2509 --omni --port 8092
 ```
 
+### JoyAI-Image-Edit
+
+```bash
+vllm serve jdopensource/JoyAI-Image-Edit-Diffusers --omni --port 8092
+```
+
 ### Start with Parameters
 
 
@@ -76,6 +86,11 @@ curl -s http://localhost:8092/v1/chat/completions \
   | cut -d',' -f2 | base64 -d > output.png
 ```
 
+For JoyAI-Image-Edit, set classifier-free guidance with `guidance_scale`
+in `extra_body` (default: `4.0`). `true_cfg_scale` remains supported as an
+alias; if both are provided, their values must match. Values at or below `1.0`
+disable CFG.
+
 ### Method 2: Using OpenAI Python SDK
 
 ```python

examples/offline_inference/image_to_image/image_edit.py

@@ -39,6 +39,20 @@ Usage (multiple images):
         --cfg-scale 4.0 \
         --guidance-scale 1.0
 
+Usage (JoyAI-Image-Edit, single image):
+    python image_edit.py \
+        --model jdopensource/JoyAI-Image-Edit-Diffusers \
+        --image input.png \
+        --prompt "Change the background to a clean studio while preserving the subject." \
+        --height 1024 \
+        --width 1024 \
+        --num-inference-steps 50 \
+        --guidance-scale 4.0 \
+        --output output_joyai_edit.png
+
+    Note: JoyAI-Image-Edit snaps requested dimensions to the nearest supported
+    Joy bucket; for square outputs, use 1024x1024.
+
 Usage (with cache-dit acceleration):
     python image_edit.py \
         --image input.png \
@@ -270,9 +284,10 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument(
         "--cfg-scale",
         type=float,
-        default=4.0,
+        default=None,
         help=(
-            "True classifier-free guidance scale (default: 4.0). Guidance scale as defined in Classifier-Free "
+            "True classifier-free guidance scale. When omitted, the pipeline uses its model-specific default. "
+            "Guidance scale as defined in Classifier-Free "
             "Diffusion Guidance. Classifier-free guidance is enabled by setting cfg_scale > 1 and providing "
             "a negative_prompt. Higher guidance scale encourages images closely linked to the text prompt, "
             "usually at the expense of lower image quality."
@@ -281,9 +296,9 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument(
         "--guidance-scale",
         type=float,
-        default=1.0,
+        default=None,
         help=(
-            "Guidance scale for guidance-distilled models (default: 1.0, disabled). "
+            "Guidance scale for guidance-distilled models. When omitted, the pipeline uses its model-specific default. "
             "Unlike classifier-free guidance (--cfg-scale), guidance-distilled models take the guidance scale "
             "directly as an input parameter. Enabled when guidance_scale > 1. Ignored when not using guidance-distilled models."
         ),
@@ -523,6 +538,12 @@ def parse_args() -> argparse.Namespace:
             "silently masking with a possibly-wrong prompt."
         ),
     )
+    parser.add_argument(
+        "--init-timeout",
+        type=int,
+        default=600,
+        help="Overall pipeline initialization timeout in seconds.",
+    )
     return parser.parse_args()
 
 
@@ -593,6 +614,7 @@ def main():
         enable_cpu_offload=args.enable_cpu_offload,
         enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
         profiler_config=args.profiler_config,
+        init_timeout=args.init_timeout,
     )
     if args.enforce_eager is not None:
         omni_kwargs["enforce_eager"] = args.enforce_eager

examples/offline_inference/image_to_image/image_to_image.md

@@ -1,6 +1,6 @@
 # Image-To-Image
 
-This example edits an input image with `Qwen/Qwen-Image-Edit` using the `image_edit.py` CLI.
+This example edits an input image with `Qwen/Qwen-Image-Edit` and `JoyAI-Image-Edit` using the `image_edit.py` CLI.
 
 ## Local CLI Usage
 
@@ -23,6 +23,27 @@ python image_edit.py \
   --cfg-scale 4.0
 ```
 
+### JoyAI-Image-Edit
+
+JoyAI-Image-Edit uses the same `image_edit.py` CLI and supports one input
+image per request. Requested dimensions are snapped to the nearest supported Joy
+bucket; for square outputs, use 1024x1024.
+
+```bash
+python image_edit.py \
+  --model jdopensource/JoyAI-Image-Edit-Diffusers \
+  --image input.png \
+  --prompt "Change the background to a clean studio while preserving the subject." \
+  --height 1024 \
+  --width 1024 \
+  --num-inference-steps 50 \
+  --guidance-scale 4.0 \
+  --output output_joyai_edit.png
+```
+
+JoyAI defaults to guidance scale `4.0`; pass `--guidance-scale` explicitly when
+you need a different value.
+
 ### Multiple Image Editing (Qwen-Image-Edit-2509)
 
 For multiple image inputs, use `Qwen/Qwen-Image-Edit-2509` or  `Qwen/Qwen-Image-Edit-2511`:
@@ -40,7 +61,7 @@ python image_edit.py \
 
 Key arguments:
 
-- `--model`: model name or path. Use `Qwen/Qwen-Image-Edit-2509` or later for multiple image support.
+- `--model`: model name or path. Use `Qwen/Qwen-Image-Edit-2509` or later for multiple image support. `jdopensource/JoyAI-Image-Edit-Diffusers` supports one input image.
 - `--image`: path(s) to the source image(s) (PNG/JPG, converted to RGB). Can specify multiple images.
 - `--prompt` / `--negative-prompt`: text description (string).
 - `--cfg-scale`: true classifier-free guidance scale (default: 4.0). Classifier-free guidance is enabled by setting cfg_scale > 1 and providing a negative_prompt. Higher guidance scale encourages images closely linked to the text prompt, usually at the expense of lower image quality.

examples/online_serving/image_to_image/README.md

@@ -2,10 +2,12 @@
 
 This example demonstrates how to deploy image-to-image models for online image editing service using vLLM-Omni.
 
-Supported models include Qwen-Image-Edit, BAGEL, and other image-to-image pipelines.
+Supported models include Qwen-Image-Edit, BAGEL, JoyAI-Image-Edit, and other image-to-image pipelines.
 
 For **multi-image** input editing, use **Qwen-Image-Edit-2509** (QwenImageEditPlusPipeline) and send multiple images in the user message content.
 
+For **JoyAI-Image-Edit**, use `jdopensource/JoyAI-Image-Edit-Diffusers` and send exactly one input image.
+
 ## Start Server
 
 ### Basic Start
@@ -29,6 +31,12 @@ vllm serve Qwen/Qwen-Image-Edit-2509 --omni --port 8092
 vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091
 ```
 
+### JoyAI-Image-Edit
+
+```bash
+vllm serve jdopensource/JoyAI-Image-Edit-Diffusers --omni --port 8092
+```
+
 ### Start with Parameters
 
 
@@ -77,6 +85,11 @@ EOF
 curl -s http://localhost:8092/v1/chat/completions   -H "Content-Type: application/json"   -d @request.json | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2 | base64 -d > output.png
 ```
 
+For JoyAI-Image-Edit, set classifier-free guidance with `guidance_scale`
+in `extra_body` (default: `4.0`). `true_cfg_scale` remains supported as an
+alias; if both are provided, their values must match. Values at or below `1.0`
+disable CFG.
+
 ### Method 2: Using OpenAI Python SDK
 
 ```python

recipes/JD/JoyAI-Image-Edit.md

@@ -0,0 +1,221 @@
+# JoyAI-Image-Edit
+
+> Text-guided single-image editing with the native JoyAI-Image-Edit diffusion pipeline
+
+## Summary
+
+- Vendor: JD / JoyImage Team
+- Model: `jdopensource/JoyAI-Image-Edit-Diffusers`
+- Task: Text-guided single-image editing
+- Mode: Offline inference and OpenAI-compatible online serving
+- Maintainer: Community
+
+## When to use this recipe
+
+Use this recipe when you want to run JoyAI-Image-Edit through the shared
+image-to-image examples with a known-good baseline configuration. JoyAI-Image-Edit
+is a large single-image editing model: it requires one reference image per
+request, snaps requested sizes to Joy's supported resolution buckets, and has a
+validated 1024x1024 baseline on an 80 GB NVIDIA GPU.
+
+## References
+
+- Upstream model:
+  [`jdopensource/JoyAI-Image-Edit-Diffusers`](https://huggingface.co/jdopensource/JoyAI-Image-Edit-Diffusers)
+- Upstream project:
+  [`jd-opensource/JoyAI-Image`](https://github.com/jd-opensource/JoyAI-Image)
+- Related offline example:
+  [`examples/offline_inference/image_to_image/image_edit.py`](../../examples/offline_inference/image_to_image/image_edit.py)
+- Related online example:
+  [`examples/online_serving/image_to_image/README.md`](../../examples/online_serving/image_to_image/README.md)
+- User guide:
+  [`docs/user_guide/examples/offline_inference/image_to_image.md`](../../docs/user_guide/examples/offline_inference/image_to_image.md),
+  [`docs/user_guide/examples/online_serving/image_to_image.md`](../../docs/user_guide/examples/online_serving/image_to_image.md)
+
+## Hardware Support
+
+This recipe documents the CUDA configuration used to validate the PR's JoyAI
+integration. Other GPU types should be treated as unvalidated until you rerun
+the smoke and parity checks below.
+
+## GPU
+
+### 1x NVIDIA H800 PCIe 80GB
+
+#### Environment
+
+- OS: Linux
+- Python: `3.12.3`
+- Driver / runtime: NVIDIA driver `580.82.07`, CUDA runtime `13.0`
+- PyTorch: `2.11.0+cu130`
+- Diffusers: `0.39.0.dev0` with `JoyImageEditPipeline` available
+- vLLM: `0.21.0`
+- vLLM-Omni: use the commit you are deploying from
+- GPU: NVIDIA H800 PCIe, 81559 MiB
+
+#### Online Serving
+
+Start the OpenAI-compatible server from the repository root:
+
+```bash
+CUDA_VISIBLE_DEVICES=0 \
+PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \
+vllm serve jdopensource/JoyAI-Image-Edit-Diffusers \
+  --omni \
+  --port 8092 \
+  --init-timeout 1200
+```
+
+The server is ready when the log shows `Application startup complete.`
+
+If you need extra GPU headroom, use one of the memory-saving variants below.
+These variants are useful on tighter 80 GB deployments, but the latency and
+quality numbers in this recipe are from the baseline command above.
+
+```bash
+# Pipeline-level CPU offload for large components.
+CUDA_VISIBLE_DEVICES=0 \
+PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \
+vllm serve jdopensource/JoyAI-Image-Edit-Diffusers \
+  --omni \
+  --port 8092 \
+  --init-timeout 1200 \
+  --enable-cpu-offload
+```
+
+```bash
+# DiT blockwise offload. This trades latency for lower peak GPU memory.
+CUDA_VISIBLE_DEVICES=0 \
+PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \
+vllm serve jdopensource/JoyAI-Image-Edit-Diffusers \
+  --omni \
+  --port 8092 \
+  --init-timeout 1200 \
+  --enable-layerwise-offload
+```
+
+#### Online Verification
+
+Download the public JoyAI plate example:
+
+```bash
+wget -O /tmp/joy_plate.jpg \
+  https://raw.githubusercontent.com/jd-opensource/JoyAI-Image/main/test_images/test_1.jpg
+```
+
+Send one image and one prompt. JoyAI-Image-Edit supports exactly one input image
+per request.
+
+```bash
+IMG_B64=$(base64 -w0 /tmp/joy_plate.jpg)
+
+curl -s http://localhost:8092/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d "$(jq -n --arg img "$IMG_B64" '{
+    messages: [{
+      role: "user",
+      content: [
+        {type: "text", text: "Turn the plate blue"},
+        {type: "image_url", image_url: {url: ("data:image/jpeg;base64," + $img)}}
+      ]
+    }],
+    extra_body: {
+      height: 1024,
+      width: 1024,
+      num_inference_steps: 30,
+      guidance_scale: 4.0,
+      seed: 0
+    }
+  }')" \
+  | jq -r '.choices[0].message.content[0].image_url.url' \
+  | cut -d',' -f2- \
+  | base64 -d > /tmp/joyai_image_edit_plate_online.png
+```
+
+Expected result: `/tmp/joyai_image_edit_plate_online.png` is a 1024x1024 PNG
+where the plate color changes to blue while the surrounding scene is mostly
+preserved. Use a client timeout of at least 600 seconds for first-run or cold
+server checks.
+
+#### Offline Command
+
+Run the same smoke case without starting a server:
+
+```bash
+CUDA_VISIBLE_DEVICES=0 \
+PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \
+python examples/offline_inference/image_to_image/image_edit.py \
+  --model jdopensource/JoyAI-Image-Edit-Diffusers \
+  --image /tmp/joy_plate.jpg \
+  --prompt "Turn the plate blue" \
+  --negative-prompt "" \
+  --height 1024 \
+  --width 1024 \
+  --num-inference-steps 30 \
+  --guidance-scale 4.0 \
+  --seed 0 \
+  --enforce-eager \
+  --init-timeout 1200 \
+  --output /tmp/joyai_image_edit_plate_offline.png
+```
+
+For a local checkpoint mirror, replace `--model` with the local model directory.
+The PR validation used the same command shape with
+`/autodl-fs/data/joyai_models/JoyAI-Image-Edit-Diffusers`.
+
+#### Validated Results
+
+The table below is from the 1x H800 PCIe environment above. Resolution requests
+were snapped to the Joy 1024x1024 bucket, batch size was 1, and CFG was 4.0.
+`vLLM peak` is the observed peak GPU memory for the vLLM-Omni run.
+
+| Case | Steps | Seed | vLLM latency | Diffusers latency | vLLM peak | Diffusers peak | SSIM | PSNR |
+| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: |
+| `plate_seed123` | 30 | 123 | 50.87 s | 52.74 s | 57.52 GiB | 57.28 GiB | 0.9869 | 34.7051 |
+| `plate_seed0` | 30 | 0 | 51.02 s | 52.01 s | 57.52 GiB | 57.28 GiB | 0.9946 | 41.4351 |
+| `astronaut_seed0` | 40 | 0 | 67.90 s | 69.22 s | 57.52 GiB | 57.28 GiB | 0.9948 | 38.2498 |
+| `crane_removal_seed0` | 40 | 0 | 68.00 s | 69.27 s | 57.52 GiB | 57.28 GiB | 0.9975 | 50.0135 |
+| `object_move_seed0` | 40 | 0 | 68.06 s | 69.32 s | 57.52 GiB | 57.28 GiB | 0.9969 | 46.1138 |
+| `object_rotation_seed0` | 40 | 0 | 67.94 s | 69.28 s | 57.52 GiB | 57.28 GiB | 0.9898 | 35.5365 |
+| `camera_control_seed0` | 40 | 0 | 68.15 s | 69.51 s | 57.52 GiB | 57.28 GiB | 0.9952 | 41.7697 |
+
+#### Output Size Behavior
+
+JoyAI-Image-Edit does not use arbitrary image sizes directly. The pipeline maps
+the input or requested `height` / `width` to the nearest supported Joy bucket,
+then resizes and center-crops the reference image to that bucket before
+denoising.
+
+- Provide both `height` and `width`, or omit both.
+- For square outputs, request `height=1024` and `width=1024`.
+- A request such as `512x512` is also square and is snapped to the 1024x1024 Joy
+  bucket, so the final image size will not match the literal request.
+- Non-square requests are snapped to the nearest aspect-ratio bucket supported
+  by the model.
+
+#### Guidance Parameters
+
+Use `guidance_scale` for JoyAI classifier-free guidance, matching Diffusers.
+The pipeline defaults to `4.0` when the parameter is omitted. `true_cfg_scale`
+remains supported as an alias; if both parameters are supplied, their values
+must match. Values greater than `1.0` enable CFG; an omitted negative prompt is
+treated as an empty string. Values at or below `1.0` disable CFG.
+
+The offline example passes `--guidance-scale 4.0` explicitly for clarity; if it
+is omitted, the JoyAI pipeline uses the same `4.0` default.
+
+#### Known Limitations
+
+- Single image only: JoyAI-Image-Edit rejects requests with zero or multiple
+  input images.
+- Single prompt only: batched prompt requests are rejected.
+- HSDP is not supported yet. Starting with HSDP enabled raises a `ValueError`
+  before the generic HSDP setup reaches the transformer.
+- Tensor Parallelism is not supported yet. Keep `tensor_parallel_size=1`;
+  larger values raise a `ValueError` during pipeline initialization.
+- LoRA is not included for this JoyAI pipeline.
+- Cache-DiT and TeaCache are not included for this JoyAI pipeline. Do not treat
+  the generic image-to-image cache flags as validated JoyAI acceleration knobs.
+- Text encoder quantization is not included.
+- The validation above is CUDA-only; ROCm and NPU configurations are not
+  documented as supported by this recipe.

recipes/README.md

@@ -67,6 +67,7 @@ recipes/
 | [`Kandinsky/Kandinsky6-TI2VA.md`](./Kandinsky/Kandinsky6-TI2VA.md) | Text/image-to-video-and-audio, offline and `/v1/videos` | 1x H100 80GB with CPU offload; 2x H100 for parallel modes |
 | [`LTX/LTX-2.md`](./LTX/LTX-2.md) | LTX-2/LTX-2.3 text-to-video and image-to-video with synchronized audio | H200 141GB / 96GB-class GPU / 1x Intel Arc Pro B70 32GB |
 | [`LTX/LTX-2.5.md`](./LTX/LTX-2.5.md) | LTX-2.5-Diffusers: Full/SFT one-stage and distilled two-stage T2V/I2V with synchronized audio | NVIDIA B300; cuDNN-qualified / 1x Intel Arc Pro B70 32GB |
+| [`JD/JoyAI-Image-Edit.md`](./JD/JoyAI-Image-Edit.md) | Text-guided single-image editing | 1x H800 PCIe 80GB |
 | [`m-a-p/SheetSage2-H200.md`](./m-a-p/SheetSage2-H200.md) | Official audio-to-score preprocessing and YuE2 request export | 1x H200 141GB |
 | [`m-a-p/YuE2-3B.md`](./m-a-p/YuE2-3B.md) | Offline and online text-to-music with optional ABC score | 1x RTX 4090 24GB |
 | [`MammothModa2/MammothModa2.md`](./MammothModa2/MammothModa2.md) | Preview and Dev text-to-image (AR → DiT); Dev text/image understanding | Preview: 1x L40S 48GB / 1x ≥40GB GPU; Dev: 1x NVIDIA GPU with sufficient cache headroom |

tests/e2e/accuracy/test_joy_image_edit.py

@@ -0,0 +1,401 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import gc
+import os
+from pathlib import Path
+from typing import Any
+
+import pytest
+import requests
+import torch
+from PIL import Image
+from vllm.config.load import LoadConfig
+
+from benchmarks.accuracy.common import decode_base64_image, pil_to_png_bytes
+from tests.e2e.accuracy.helpers import assert_similarity, model_output_dir
+from tests.helpers.clean import cleanup_test_environment
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniServer
+from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+pytestmark = [pytest.mark.full_model, pytest.mark.diffusion]
+
+
+MODEL_ID = "jdopensource/JoyAI-Image-Edit-Diffusers"
+MODEL_ENV_VAR = "JOY_IMAGE_EDIT_MODEL"
+INPUT_IMAGE_PATH = Path(__file__).resolve().parents[3] / "docs" / "assets" / "WeChat.jpg"
+PROMPT = "Change the background to a clean studio while preserving the subject."
+NEGATIVE_PROMPT = ""
+WIDTH = 1024
+HEIGHT = 1024
+NUM_INFERENCE_STEPS = 20
+GUIDANCE_SCALE = 4.0
+SEED = 42
+SSIM_THRESHOLD = 0.90
+PSNR_THRESHOLD = 25.0
+SERVER_ARGS = [
+    "--num-gpus",
+    "1",
+    "--enforce-eager",
+    "--init-timeout",
+    "1200",
+]
+
+
+def _model_name() -> str:
+    return os.environ.get(MODEL_ENV_VAR, MODEL_ID)
+
+
+def _local_files_only(model: str) -> bool:
+    return Path(model).exists()
+
+
+def _joy_image_edit_pipeline_cls() -> type[Any]:
+    diffusers = pytest.importorskip("diffusers")
+    pipeline_cls = getattr(diffusers, "JoyImageEditPipeline", None)
+    if pipeline_cls is None:
+        pytest.skip("diffusers.JoyImageEditPipeline is required for JoyAI-Image-Edit parity testing.")
+    return pipeline_cls
+
+
+def _load_input_image() -> Image.Image:
+    if not INPUT_IMAGE_PATH.exists():
+        raise AssertionError(f"JoyAI input image does not exist: {INPUT_IMAGE_PATH}")
+    image = Image.open(INPUT_IMAGE_PATH).convert("RGB")
+    image.load()
+    return image
+
+
+def _latent_metrics(
+    vllm_latents: torch.Tensor,
+    diffusers_latents: torch.Tensor,
+) -> dict[str, float]:
+    vllm_flat = vllm_latents.detach().float().cpu().flatten()
+    diffusers_flat = diffusers_latents.detach().float().cpu().flatten()
+    if vllm_flat.shape != diffusers_flat.shape:
+        raise AssertionError(
+            "Latent outputs have different shapes: "
+            f"vllm_omni={tuple(vllm_latents.shape)}, diffusers={tuple(diffusers_latents.shape)}"
+        )
+    delta = vllm_flat - diffusers_flat
+    return {
+        "max_abs": delta.abs().max().item(),
+        "mean_abs": delta.abs().mean().item(),
+        "rmse": torch.sqrt(torch.mean(delta.square())).item(),
+        "cosine_similarity": torch.nn.functional.cosine_similarity(
+            vllm_flat.unsqueeze(0),
+            diffusers_flat.unsqueeze(0),
+        ).item(),
+    }
+
+
+def _print_latent_metrics(metrics: dict[str, float]) -> None:
+    print(f"{MODEL_ID} latent metrics:")
+    for name, value in metrics.items():
+        print(f"  {name}: {value:.8f}")
+
+
+def _ensure_diffusion_distributed_initialized() -> None:
+    from vllm_omni.diffusion.distributed.parallel_state import (
+        init_distributed_environment,
+        initialize_model_parallel,
+        model_parallel_is_initialized,
+    )
+
+    if not torch.distributed.is_initialized():
+        os.environ.setdefault("MASTER_ADDR", "localhost")
+        os.environ.setdefault("MASTER_PORT", "29513")
+        os.environ.setdefault("RANK", "0")
+        os.environ.setdefault("LOCAL_RANK", "0")
+        os.environ.setdefault("WORLD_SIZE", "1")
+        init_distributed_environment(world_size=1, rank=0, local_rank=0)
+    if not model_parallel_is_initialized():
+        initialize_model_parallel()
+
+
+def _run_vllm_omni_joy_image_edit(
+    *,
+    model: str,
+    input_image: Image.Image,
+    output_path: Path,
+) -> Image.Image:
+    with OmniServer(model=model, serve_args=SERVER_ARGS) as server:
+        response = requests.post(
+            f"http://{server.host}:{server.port}/v1/images/edits",
+            data={
+                "model": server.model,
+                "prompt": PROMPT,
+                "size": f"{WIDTH}x{HEIGHT}",
+                "n": 1,
+                "response_format": "b64_json",
+                "negative_prompt": NEGATIVE_PROMPT,
+                "num_inference_steps": NUM_INFERENCE_STEPS,
+                "guidance_scale": GUIDANCE_SCALE,
+                "seed": SEED,
+            },
+            files=[("image", ("input.png", pil_to_png_bytes(input_image), "image/png"))],
+            timeout=1200,
+        )
+        response.raise_for_status()
+        payload = response.json()
+        assert len(payload["data"]) == 1
+        image = decode_base64_image(payload["data"][0]["b64_json"])
+        image.load()
+        image.save(output_path)
+        return image
+
+
+class _JoyImageEditRequest:
+    def __init__(self, input_image: Image.Image) -> None:
+        self.prompts = [
+            {
+                "prompt": PROMPT,
+                "negative_prompt": NEGATIVE_PROMPT,
+                "multi_modal_data": {"image": input_image},
+            }
+        ]
+        self.sampling_params = OmniDiffusionSamplingParams(
+            height=HEIGHT,
+            width=WIDTH,
+            num_inference_steps=NUM_INFERENCE_STEPS,
+            guidance_scale=GUIDANCE_SCALE,
+            guidance_scale_provided=True,
+            seed=SEED,
+        )
+        self.request_id = "joy-image-edit-latent-parity"
+
+    def is_dummy_run(self) -> bool:
+        return False
+
+
+def _load_vllm_omni_joy_image_edit_pipeline(model: str):
+    from vllm_omni.diffusion.models.joy_image.pipeline_joy_image_edit import (
+        JoyImageEditPipeline,
+    )
+
+    _ensure_diffusion_distributed_initialized()
+    od_config = OmniDiffusionConfig(
+        model=model,
+        model_class_name="JoyImageEditPipeline",
+        dtype=torch.bfloat16,
+        enforce_eager=True,
+    )
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+    return loader.load_model(
+        load_device="cuda",
+        load_format="custom_pipeline",
+        custom_pipeline_name=JoyImageEditPipeline,
+    )
+
+
+def _extract_diffusers_latents(result: Any) -> torch.Tensor:
+    if isinstance(result, torch.Tensor):
+        return result
+    images = getattr(result, "images", None)
+    if isinstance(images, torch.Tensor):
+        return images
+    raise TypeError(f"Unexpected Diffusers latent output type: {type(result)!r}")
+
+
+def _run_latent_parity_comparison(
+    *,
+    model: str,
+    input_image: Image.Image,
+) -> dict[str, float]:
+    pipeline_cls = _joy_image_edit_pipeline_cls()
+    cleanup_test_environment()
+    diffusers_pipe = None
+    try:
+        diffusers_pipe = pipeline_cls.from_pretrained(
+            model,
+            torch_dtype=torch.bfloat16,
+            trust_remote_code=True,
+            local_files_only=_local_files_only(model),
+        ).to("cuda")
+        if hasattr(diffusers_pipe, "set_progress_bar_config"):
+            diffusers_pipe.set_progress_bar_config(disable=False)
+
+        processed_image = diffusers_pipe.vae_image_processor.resize_center_crop(
+            input_image,
+            (HEIGHT, WIDTH),
+        )
+        prompt_embeds, prompt_embeds_mask = diffusers_pipe.encode_prompt_multiple_images(
+            prompt=PROMPT,
+            images=processed_image,
+            device=torch.device("cuda"),
+            num_images_per_prompt=1,
+            max_sequence_length=4096,
+        )
+        negative_prompt_embeds, negative_prompt_embeds_mask = diffusers_pipe.encode_prompt_multiple_images(
+            prompt=NEGATIVE_PROMPT,
+            images=processed_image,
+            device=torch.device("cuda"),
+            num_images_per_prompt=1,
+            max_sequence_length=4096,
+        )
+        num_channels_latents = int(diffusers_pipe.transformer.config.in_channels)
+        generator = torch.Generator(device="cuda").manual_seed(SEED)
+        fixed_latents = torch.randn(
+            1,
+            1,
+            num_channels_latents,
+            1,
+            HEIGHT // diffusers_pipe.vae_scale_factor_spatial,
+            WIDTH // diffusers_pipe.vae_scale_factor_spatial,
+            generator=generator,
+            device="cuda",
+            dtype=prompt_embeds.dtype,
+        )
+        diffusers_result = diffusers_pipe(
+            prompt=None,
+            image=input_image,
+            height=HEIGHT,
+            width=WIDTH,
+            num_inference_steps=NUM_INFERENCE_STEPS,
+            guidance_scale=GUIDANCE_SCALE,
+            latents=fixed_latents,
+            prompt_embeds=prompt_embeds,
+            prompt_embeds_mask=prompt_embeds_mask,
+            negative_prompt_embeds=negative_prompt_embeds,
+            negative_prompt_embeds_mask=negative_prompt_embeds_mask,
+            output_type="latent",
+        )
+        diffusers_latents = _extract_diffusers_latents(diffusers_result).detach().cpu()
+        prompt_embeds = prompt_embeds.detach().cpu()
+        prompt_embeds_mask = prompt_embeds_mask.detach().cpu()
+        negative_prompt_embeds = negative_prompt_embeds.detach().cpu()
+        negative_prompt_embeds_mask = negative_prompt_embeds_mask.detach().cpu()
+        fixed_latents = fixed_latents.detach().cpu()
+    finally:
+        if diffusers_pipe is not None and hasattr(diffusers_pipe, "maybe_free_model_hooks"):
+            diffusers_pipe.maybe_free_model_hooks()
+        del diffusers_pipe
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.accelerator.empty_cache()
+        cleanup_test_environment()
+
+    cleanup_test_environment()
+    vllm_pipe = None
+    try:
+        from vllm_omni.diffusion.models.joy_image.pipeline_joy_image_edit import (
+            get_joy_image_edit_pre_process_func,
+        )
+
+        vllm_pipe = _load_vllm_omni_joy_image_edit_pipeline(model)
+        vllm_pipe.eval()
+        request = _JoyImageEditRequest(input_image)
+        request = get_joy_image_edit_pre_process_func(vllm_pipe.od_config)(request)
+        with torch.inference_mode():
+            vllm_result = vllm_pipe(
+                request,
+                prompt_embeds=prompt_embeds.to(device=vllm_pipe.device, dtype=torch.bfloat16),
+                prompt_embeds_mask=prompt_embeds_mask.to(device=vllm_pipe.device),
+                negative_prompt_embeds=negative_prompt_embeds.to(device=vllm_pipe.device, dtype=torch.bfloat16),
+                negative_prompt_embeds_mask=negative_prompt_embeds_mask.to(device=vllm_pipe.device),
+                latents=fixed_latents.to(device=vllm_pipe.device, dtype=torch.bfloat16),
+                output_type="latent",
+            )
+        vllm_latents = vllm_result.output
+    finally:
+        if vllm_pipe is not None and hasattr(vllm_pipe, "maybe_free_model_hooks"):
+            vllm_pipe.maybe_free_model_hooks()
+        del vllm_pipe
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.accelerator.empty_cache()
+        cleanup_test_environment()
+
+    return _latent_metrics(vllm_latents, diffusers_latents)
+
+
+def _run_diffusers_joy_image_edit(
+    *,
+    model: str,
+    input_image: Image.Image,
+    output_path: Path,
+) -> Image.Image:
+    pipeline_cls = _joy_image_edit_pipeline_cls()
+    cleanup_test_environment()
+    pipe = None
+    try:
+        pipe = pipeline_cls.from_pretrained(
+            model,
+            torch_dtype=torch.bfloat16,
+            trust_remote_code=True,
+            local_files_only=_local_files_only(model),
+        ).to("cuda")
+        if hasattr(pipe, "set_progress_bar_config"):
+            pipe.set_progress_bar_config(disable=False)
+        generator = torch.Generator(device="cuda").manual_seed(SEED)
+        result = pipe(
+            prompt=PROMPT,
+            image=input_image,
+            negative_prompt=NEGATIVE_PROMPT,
+            height=HEIGHT,
+            width=WIDTH,
+            num_inference_steps=NUM_INFERENCE_STEPS,
+            guidance_scale=GUIDANCE_SCALE,
+            generator=generator,
+        )
+        output_image = result.images[0].convert("RGB")
+        output_image.save(output_path)
+        return output_image
+    finally:
+        if pipe is not None and hasattr(pipe, "maybe_free_model_hooks"):
+            pipe.maybe_free_model_hooks()
+        del pipe
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.accelerator.empty_cache()
+        cleanup_test_environment()
+
+
+@pytest.mark.benchmark
+@hardware_test(res={"cuda": ["H100", "B200"]}, num_cards=1)
+def test_joy_image_edit_matches_diffusers(accuracy_artifact_root: Path) -> None:
+    _joy_image_edit_pipeline_cls()
+    model = _model_name()
+    output_dir = model_output_dir(accuracy_artifact_root, MODEL_ID)
+    input_image = _load_input_image()
+    input_path = output_dir / "input.png"
+    vllm_output_path = output_dir / "vllm_omni.png"
+    diffusers_output_path = output_dir / "diffusers.png"
+    input_image.save(input_path)
+
+    vllm_image = _run_vllm_omni_joy_image_edit(
+        model=model,
+        input_image=input_image,
+        output_path=vllm_output_path,
+    )
+    diffusers_image = _run_diffusers_joy_image_edit(
+        model=model,
+        input_image=input_image,
+        output_path=diffusers_output_path,
+    )
+
+    print(f"{MODEL_ID} generated images:")
+    print(f"  input: {input_path}")
+    print(f"  vllm_omni: {vllm_output_path}")
+    print(f"  diffusers: {diffusers_output_path}")
+
+    assert_similarity(
+        model_name=MODEL_ID,
+        vllm_image=vllm_image,
+        diffusers_image=diffusers_image,
+        width=WIDTH,
+        height=HEIGHT,
+        ssim_threshold=SSIM_THRESHOLD,
+        psnr_threshold=PSNR_THRESHOLD,
+    )
+    _print_latent_metrics(
+        _run_latent_parity_comparison(
+            model=model,
+            input_image=input_image,
+        )
+    )

tests/model_tests/diffusion/test_alignment.py

@@ -25,6 +25,9 @@ pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
 # for adding a new tiny model builder and corresponding entry in DIFFUSION_TEST_SETTINGS.
 EXCLUDED_MODELS = [
     "QwenImageLayeredPipeline",
+    # Uses a custom multi-image editing contract; covered by the dedicated
+    # JoyImage tiny-pipeline tests and real-checkpoint E2E accuracy test.
+    "JoyImageEditPipeline",
     # Dedicated CPU/unit coverage in tests/diffusion/models/qwen_image_21/.
     # Tiny-model offline suite is not wired yet (2.1 transformer/VAE configs).
     "QwenImage21Pipeline",

vllm_omni/diffusion/model_metadata.py

@@ -30,6 +30,7 @@ QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES = 4
 QWEN_IMAGE_21_MAX_INPUT_IMAGES = 4
 # Upstream HunyuanImage-3.0 "Multi-Image Fusion" caps reference images at 3.
 HUNYUAN_IMAGE3_MAX_INPUT_IMAGES = 3
+JOY_IMAGE_EDIT_MAX_INPUT_IMAGES = 1
 # Boogu-Image editing (TI2I) supports a single reference image for now.
 BOOGU_IMAGE_MAX_INPUT_IMAGES = 1
 
@@ -61,6 +62,10 @@ _DIFFUSION_MODEL_METADATA: dict[str, DiffusionModelMetadata] = {
         supports_multimodal_inputs=True,
         max_multimodal_image_inputs=BOOGU_IMAGE_MAX_INPUT_IMAGES,
     ),
+    "JoyImageEditPipeline": DiffusionModelMetadata(
+        supports_multimodal_inputs=True,
+        max_multimodal_image_inputs=JOY_IMAGE_EDIT_MAX_INPUT_IMAGES,
+    ),
     "MiniMaxH3Pipeline": DiffusionModelMetadata(
         supports_multimodal_inputs=True,
         max_multimodal_image_inputs=9,

vllm_omni/diffusion/models/joy_image/__init__.py

@@ -0,0 +1,18 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from vllm_omni.diffusion.models.joy_image.joy_image_edit_transformer import (
+    JoyImageEditTransformer3DModel,
+)
+from vllm_omni.diffusion.models.joy_image.pipeline_joy_image_edit import (
+    JoyImageEditPipeline,
+    get_joy_image_edit_post_process_func,
+    get_joy_image_edit_pre_process_func,
+)
+
+__all__ = [
+    "JoyImageEditPipeline",
+    "JoyImageEditTransformer3DModel",
+    "get_joy_image_edit_pre_process_func",
+    "get_joy_image_edit_post_process_func",
+]

vllm_omni/diffusion/models/joy_image/cfg_parallel.py

@@ -0,0 +1,83 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import logging
+
+import torch
+
+from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
+from vllm_omni.diffusion.distributed.parallel_state import (
+    get_classifier_free_guidance_world_size,
+)
+from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
+
+logger = logging.getLogger(__name__)
+
+
+class JoyImageEditCFGParallelMixin(CFGParallelMixin, ProgressBarMixin):
+    def cfg_normalize_function(self, noise_pred: torch.Tensor, comb_pred: torch.Tensor) -> torch.Tensor:
+        cond_norm = torch.norm(noise_pred, dim=2, keepdim=True)
+        comb_norm = torch.norm(comb_pred, dim=2, keepdim=True)
+        return comb_pred * (cond_norm / comb_norm.clamp_min(1e-6))
+
+    def diffuse(
+        self,
+        latents: torch.Tensor,
+        image_latents: torch.Tensor,
+        prompt_embeds: torch.Tensor,
+        prompt_embeds_mask: torch.Tensor,
+        negative_prompt_embeds: torch.Tensor | None,
+        negative_prompt_embeds_mask: torch.Tensor | None,
+        timesteps: torch.Tensor,
+        do_true_cfg: bool,
+        true_cfg_scale: float,
+        cfg_normalize: bool = True,
+    ) -> torch.Tensor:
+        self.scheduler.set_begin_index(0)
+        with self.progress_bar(total=len(timesteps)) as pbar:
+            for timestep in timesteps:
+                if self.interrupt:
+                    continue
+                self._current_timestep = timestep
+
+                latents = torch.cat([image_latents, latents[:, image_latents.shape[1] :]], dim=1)
+                latent_model_input = latents
+                timestep_expand = timestep.expand(latents.shape[0]).to(device=latents.device)
+
+                positive_kwargs = {
+                    "hidden_states": latent_model_input,
+                    "timestep": timestep_expand,
+                    "encoder_hidden_states": prompt_embeds,
+                    "encoder_hidden_states_mask": prompt_embeds_mask,
+                    "return_dict": False,
+                }
+                negative_kwargs = None
+                if do_true_cfg:
+                    negative_kwargs = {
+                        "hidden_states": latent_model_input,
+                        "timestep": timestep_expand,
+                        "encoder_hidden_states": negative_prompt_embeds,
+                        "encoder_hidden_states_mask": negative_prompt_embeds_mask,
+                        "return_dict": False,
+                    }
+
+                noise_pred = self.predict_noise_maybe_with_cfg(
+                    do_true_cfg=do_true_cfg,
+                    true_cfg_scale=true_cfg_scale,
+                    positive_kwargs=positive_kwargs,
+                    negative_kwargs=negative_kwargs,
+                    cfg_normalize=cfg_normalize,
+                    output_slice=None,
+                )
+                latents = self.scheduler_step_maybe_with_cfg(noise_pred, timestep, latents, do_true_cfg)
+                latents = torch.cat([image_latents, latents[:, image_latents.shape[1] :]], dim=1)
+                pbar.update()
+        return latents
+
+    def check_cfg_parallel_validity(self, true_cfg_scale: float) -> bool:
+        if get_classifier_free_guidance_world_size() == 1:
+            return True
+        if true_cfg_scale <= 1:
+            logger.warning("CFG parallel is enabled but Joy true_cfg_scale <= 1, so only the positive branch is used.")
+            return False
+        return True

vllm_omni/diffusion/models/joy_image/joy_image_edit_transformer.py

@@ -0,0 +1,692 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+# SPDX-FileCopyrightText: Copyright 2025 The JoyImage Team and The HuggingFace Team
+# Adapted from Hugging Face Diffusers commit
+# 23ba73e1d2079c4b89959484ed0ca1c22e7ef998:
+# src/diffusers/models/transformers/transformer_joyimage.py
+
+from __future__ import annotations
+
+import json
+import math
+from collections.abc import Iterable
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from vllm.logger import init_logger
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+
+from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+from vllm_omni.diffusion.attention.layer import Attention
+
+if TYPE_CHECKING:
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+logger = init_logger(__name__)
+
+
+def _raise_if_unsupported_tensor_parallel(od_config: OmniDiffusionConfig | None) -> None:
+    parallel_config = getattr(od_config, "parallel_config", None)
+    tensor_parallel_size = getattr(parallel_config, "tensor_parallel_size", 1) or 1
+    if tensor_parallel_size > 1:
+        raise ValueError(
+            "JoyImageEditPipeline and JoyImageEditTransformer3DModel do not support "
+            "Tensor Parallelism yet. "
+            "Please set `tensor_parallel_size=1`, or implement TP-aware linear layers "
+            "and TP parity tests."
+        )
+
+
+def _as_3tuple(value: int | list[int] | tuple[int, int, int]) -> tuple[int, int, int]:
+    if isinstance(value, int):
+        return (value, value, value)
+    if len(value) != 3:
+        raise ValueError(f"Expected a 3D patch size, got {value!r}.")
+    return (int(value[0]), int(value[1]), int(value[2]))
+
+
+def _rotate_half(x: torch.Tensor) -> torch.Tensor:
+    x_real, x_imag = x.float().reshape(*x.shape[:-1], -1, 2).unbind(-1)
+    return torch.stack([-x_imag, x_real], dim=-1).flatten(3)
+
+
+def _apply_rotary_emb(
+    xq: torch.Tensor,
+    xk: torch.Tensor,
+    freqs_cis: tuple[torch.Tensor, torch.Tensor],
+) -> tuple[torch.Tensor, torch.Tensor]:
+    ndim = xq.ndim
+    shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(xq.shape)]
+    cos = freqs_cis[0].view(*shape).to(xq.device)
+    sin = freqs_cis[1].view(*shape).to(xq.device)
+    xq_out = (xq.float() * cos + _rotate_half(xq) * sin).type_as(xq)
+    xk_out = (xk.float() * cos + _rotate_half(xk) * sin).type_as(xk)
+    return xq_out, xk_out
+
+
+class FP32LayerNorm(nn.LayerNorm):
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        dtype = hidden_states.dtype
+        hidden_states = F.layer_norm(
+            hidden_states.float(),
+            self.normalized_shape,
+            self.weight.float() if self.weight is not None else None,
+            self.bias.float() if self.bias is not None else None,
+            self.eps,
+        )
+        return hidden_states.to(dtype)
+
+
+class GeluApproximate(nn.Module):
+    def __init__(self, dim: int, inner_dim: int) -> None:
+        super().__init__()
+        self.proj = nn.Linear(dim, inner_dim)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return F.gelu(self.proj(hidden_states), approximate="tanh")
+
+
+class FeedForward(nn.Module):
+    """Small local equivalent of the Diffusers FeedForward used by JoyImage."""
+
+    def __init__(self, dim: int, inner_dim: int) -> None:
+        super().__init__()
+        self.net = nn.ModuleList(
+            [
+                GeluApproximate(dim, inner_dim),
+                nn.Identity(),
+                nn.Linear(inner_dim, dim),
+            ]
+        )
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        for layer in self.net:
+            hidden_states = layer(hidden_states)
+        return hidden_states
+
+
+class Timesteps(nn.Module):
+    def __init__(
+        self,
+        num_channels: int,
+        flip_sin_to_cos: bool = True,
+        downscale_freq_shift: float = 0,
+    ) -> None:
+        super().__init__()
+        self.num_channels = num_channels
+        self.flip_sin_to_cos = flip_sin_to_cos
+        self.downscale_freq_shift = downscale_freq_shift
+
+    def forward(self, timesteps: torch.Tensor) -> torch.Tensor:
+        timesteps = timesteps.flatten().float()
+        half_dim = self.num_channels // 2
+        exponent = -math.log(10000) * torch.arange(half_dim, dtype=torch.float32, device=timesteps.device)
+        exponent = exponent / (half_dim - self.downscale_freq_shift)
+        emb = timesteps[:, None] * torch.exp(exponent)[None, :]
+        emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)
+        if self.flip_sin_to_cos:
+            emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1)
+        if self.num_channels % 2 == 1:
+            emb = F.pad(emb, (0, 1))
+        return emb
+
+
+class TimestepEmbedding(nn.Module):
+    def __init__(self, in_channels: int, time_embed_dim: int) -> None:
+        super().__init__()
+        self.linear_1 = nn.Linear(in_channels, time_embed_dim)
+        self.act = nn.SiLU()
+        self.linear_2 = nn.Linear(time_embed_dim, time_embed_dim)
+
+    def forward(self, sample: torch.Tensor) -> torch.Tensor:
+        sample = self.linear_1(sample)
+        sample = self.act(sample)
+        return self.linear_2(sample)
+
+
+class PixArtAlphaTextProjection(nn.Module):
+    def __init__(self, in_features: int, hidden_size: int) -> None:
+        super().__init__()
+        self.linear_1 = nn.Linear(in_features, hidden_size)
+        self.act_1 = nn.GELU(approximate="tanh")
+        self.linear_2 = nn.Linear(hidden_size, hidden_size)
+
+    def forward(self, caption: torch.Tensor) -> torch.Tensor:
+        caption = self.linear_1(caption)
+        caption = self.act_1(caption)
+        return self.linear_2(caption)
+
+
+class JoyImageModulate(nn.Module):
+    def __init__(self, hidden_size: int, factor: int) -> None:
+        super().__init__()
+        self.factor = factor
+        self.modulate_table = nn.Parameter(torch.zeros(1, factor, hidden_size))
+
+    def forward(self, hidden_states: torch.Tensor) -> list[torch.Tensor]:
+        if hidden_states.ndim != 3:
+            hidden_states = hidden_states.unsqueeze(1)
+        return [item.squeeze(1) for item in (self.modulate_table + hidden_states).chunk(self.factor, dim=1)]
+
+
+class JoyImageAttention(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_attention_heads: int,
+        attention_head_dim: int,
+        eps: float = 1e-6,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.heads = num_attention_heads
+        self.head_dim = attention_head_dim
+        inner_dim = num_attention_heads * attention_head_dim
+
+        self.img_attn_qkv = nn.Linear(dim, inner_dim * 3, bias=True)
+        self.img_attn_q_norm = nn.RMSNorm(attention_head_dim, eps=eps)
+        self.img_attn_k_norm = nn.RMSNorm(attention_head_dim, eps=eps)
+        self.img_attn_proj = nn.Linear(inner_dim, dim, bias=True)
+
+        self.txt_attn_qkv = nn.Linear(dim, inner_dim * 3, bias=True)
+        self.txt_attn_q_norm = nn.RMSNorm(attention_head_dim, eps=eps)
+        self.txt_attn_k_norm = nn.RMSNorm(attention_head_dim, eps=eps)
+        self.txt_attn_proj = nn.Linear(inner_dim, dim, bias=True)
+
+        self.attn = Attention(
+            num_heads=self.heads,
+            head_size=self.head_dim,
+            softmax_scale=1.0 / math.sqrt(self.head_dim),
+            causal=False,
+            num_kv_heads=self.heads,
+            prefix=prefix,
+            role="joy_image.joint",
+            role_category="self",
+            qkv_layout="BSND",
+        )
+
+    def _stream_qkv(
+        self,
+        qkv_proj: nn.Linear,
+        q_norm: nn.Module,
+        k_norm: nn.Module,
+        hidden_states: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        query, key, value = qkv_proj(hidden_states).chunk(3, dim=-1)
+        query = query.unflatten(-1, (self.heads, self.head_dim))
+        key = key.unflatten(-1, (self.heads, self.head_dim))
+        value = value.unflatten(-1, (self.heads, self.head_dim))
+        return q_norm(query), k_norm(key), value
+
+    @staticmethod
+    def _normalize_attention_mask(attention_mask: torch.Tensor | None) -> torch.Tensor | None:
+        if attention_mask is None:
+            return None
+        if attention_mask.ndim == 4 and attention_mask.shape[1] == 1 and attention_mask.shape[2] == 1:
+            attention_mask = attention_mask[:, 0, 0, :]
+        if attention_mask.ndim == 2:
+            return attention_mask.to(torch.bool).contiguous()
+        return attention_mask
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        image_rotary_emb: tuple[
+            tuple[torch.Tensor, torch.Tensor] | None,
+            tuple[torch.Tensor, torch.Tensor] | None,
+        ]
+        | None = None,
+        attention_mask: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        img_query, img_key, img_value = self._stream_qkv(
+            self.img_attn_qkv,
+            self.img_attn_q_norm,
+            self.img_attn_k_norm,
+            hidden_states,
+        )
+        txt_query, txt_key, txt_value = self._stream_qkv(
+            self.txt_attn_qkv,
+            self.txt_attn_q_norm,
+            self.txt_attn_k_norm,
+            encoder_hidden_states,
+        )
+
+        if image_rotary_emb is not None:
+            vis_freqs, txt_freqs = image_rotary_emb
+            if vis_freqs is not None:
+                img_query, img_key = _apply_rotary_emb(img_query, img_key, vis_freqs)
+            if txt_freqs is not None:
+                txt_query, txt_key = _apply_rotary_emb(txt_query, txt_key, txt_freqs)
+
+        joint_query = torch.cat([img_query, txt_query], dim=1)
+        joint_key = torch.cat([img_key, txt_key], dim=1)
+        joint_value = torch.cat([img_value, txt_value], dim=1)
+
+        attn_metadata = AttentionMetadata(attn_mask=self._normalize_attention_mask(attention_mask))
+        joint_hidden_states = self.attn(joint_query, joint_key, joint_value, attn_metadata)
+        joint_hidden_states = joint_hidden_states.flatten(2, 3)
+        joint_hidden_states = joint_hidden_states.to(joint_query.dtype)
+
+        img_attn_output = joint_hidden_states[:, : hidden_states.shape[1], :]
+        txt_attn_output = joint_hidden_states[:, hidden_states.shape[1] :, :]
+        return self.img_attn_proj(img_attn_output), self.txt_attn_proj(txt_attn_output)
+
+
+class JoyImageTransformerBlock(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_attention_heads: int,
+        attention_head_dim: int,
+        mlp_width_ratio: float = 4.0,
+        eps: float = 1e-6,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        mlp_hidden_dim = int(dim * mlp_width_ratio)
+
+        self.img_mod = JoyImageModulate(dim, factor=6)
+        self.img_norm1 = FP32LayerNorm(dim, elementwise_affine=False, eps=eps)
+        self.img_norm2 = FP32LayerNorm(dim, elementwise_affine=False, eps=eps)
+        self.img_mlp = FeedForward(dim, inner_dim=mlp_hidden_dim)
+
+        self.txt_mod = JoyImageModulate(dim, factor=6)
+        self.txt_norm1 = FP32LayerNorm(dim, elementwise_affine=False, eps=eps)
+        self.txt_norm2 = FP32LayerNorm(dim, elementwise_affine=False, eps=eps)
+        self.txt_mlp = FeedForward(dim, inner_dim=mlp_hidden_dim)
+
+        self.attn = JoyImageAttention(
+            dim,
+            num_attention_heads,
+            attention_head_dim,
+            eps=eps,
+            prefix=f"{prefix}.attn" if prefix else "",
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        temb: torch.Tensor,
+        image_rotary_emb: tuple[
+            tuple[torch.Tensor, torch.Tensor] | None,
+            tuple[torch.Tensor, torch.Tensor] | None,
+        ],
+        attention_mask: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        (
+            img_mod1_shift,
+            img_mod1_scale,
+            img_mod1_gate,
+            img_mod2_shift,
+            img_mod2_scale,
+            img_mod2_gate,
+        ) = self.img_mod(temb)
+        (
+            txt_mod1_shift,
+            txt_mod1_scale,
+            txt_mod1_gate,
+            txt_mod2_shift,
+            txt_mod2_scale,
+            txt_mod2_gate,
+        ) = self.txt_mod(temb)
+
+        img_normed = self.img_norm1(hidden_states)
+        txt_normed = self.txt_norm1(encoder_hidden_states)
+        img_modulated = img_normed * (1 + img_mod1_scale.unsqueeze(1)) + img_mod1_shift.unsqueeze(1)
+        txt_modulated = txt_normed * (1 + txt_mod1_scale.unsqueeze(1)) + txt_mod1_shift.unsqueeze(1)
+        img_attn, txt_attn = self.attn(
+            hidden_states=img_modulated,
+            encoder_hidden_states=txt_modulated,
+            image_rotary_emb=image_rotary_emb,
+            attention_mask=attention_mask,
+        )
+
+        hidden_states = hidden_states + img_attn * img_mod1_gate.unsqueeze(1)
+        encoder_hidden_states = encoder_hidden_states + txt_attn * txt_mod1_gate.unsqueeze(1)
+
+        img_ffn_normed = self.img_norm2(hidden_states)
+        txt_ffn_normed = self.txt_norm2(encoder_hidden_states)
+        img_ffn_input = img_ffn_normed * (1 + img_mod2_scale.unsqueeze(1)) + img_mod2_shift.unsqueeze(1)
+        txt_ffn_input = txt_ffn_normed * (1 + txt_mod2_scale.unsqueeze(1)) + txt_mod2_shift.unsqueeze(1)
+        hidden_states = hidden_states + self.img_mlp(img_ffn_input) * img_mod2_gate.unsqueeze(1)
+        encoder_hidden_states = encoder_hidden_states + self.txt_mlp(txt_ffn_input) * txt_mod2_gate.unsqueeze(1)
+        return hidden_states, encoder_hidden_states
+
+
+class JoyImageTimeTextImageEmbedding(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        time_freq_dim: int,
+        time_proj_dim: int,
+        text_embed_dim: int,
+    ) -> None:
+        super().__init__()
+        self.timesteps_proj = Timesteps(
+            num_channels=time_freq_dim,
+            flip_sin_to_cos=True,
+            downscale_freq_shift=0,
+        )
+        self.time_embedder = TimestepEmbedding(
+            in_channels=time_freq_dim,
+            time_embed_dim=dim,
+        )
+        self.act_fn = nn.SiLU()
+        self.time_proj = nn.Linear(dim, time_proj_dim)
+        self.text_embedder = PixArtAlphaTextProjection(text_embed_dim, dim)
+
+    def forward(
+        self,
+        timestep: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        timestep = self.timesteps_proj(timestep)
+        time_embedder_dtype = next(iter(self.time_embedder.parameters())).dtype
+        if timestep.dtype != time_embedder_dtype and time_embedder_dtype != torch.int8:
+            timestep = timestep.to(time_embedder_dtype)
+        temb = self.time_embedder(timestep).type_as(encoder_hidden_states)
+        timestep_proj = self.time_proj(self.act_fn(temb))
+        encoder_hidden_states = self.text_embedder(encoder_hidden_states)
+        return temb, timestep_proj, encoder_hidden_states
+
+
+class JoyImageEditTransformer3DModel(nn.Module):
+    _skip_layerwise_casting_patterns = ["img_in", "condition_embedder", "norm"]
+    _no_split_modules = ["JoyImageTransformerBlock"]
+    _repeated_blocks = ["JoyImageTransformerBlock"]
+    _layerwise_offload_blocks_attrs = ["double_blocks"]
+
+    def __init__(
+        self,
+        *,
+        od_config: OmniDiffusionConfig | None = None,
+        patch_size: int | list[int] | tuple[int, int, int] = (1, 2, 2),
+        in_channels: int = 16,
+        out_channels: int | None = None,
+        hidden_size: int = 4096,
+        num_attention_heads: int = 32,
+        text_dim: int = 4096,
+        mlp_width_ratio: float = 4.0,
+        mlp_ratio: float | None = None,
+        num_layers: int = 40,
+        rope_dim_list: list[int] | tuple[int, ...] | None = None,
+        rope_type: str = "rope",
+        theta: int = 256,
+        **_: Any,
+    ) -> None:
+        super().__init__()
+        _raise_if_unsupported_tensor_parallel(od_config)
+        self.od_config = od_config
+        self.in_channels = in_channels
+        self.out_channels = out_channels or in_channels
+        self.hidden_size = hidden_size
+        self.text_dim = text_dim
+        self.num_layers = num_layers
+        self.num_attention_heads = num_attention_heads
+        self.patch_size = _as_3tuple(patch_size)
+        self.rope_type = rope_type
+        self.theta = theta
+
+        if hidden_size % num_attention_heads != 0:
+            raise ValueError(
+                f"hidden_size ({hidden_size}) must be divisible by num_attention_heads ({num_attention_heads})."
+            )
+        attention_head_dim = hidden_size // num_attention_heads
+        if rope_dim_list is None:
+            if attention_head_dim == 128:
+                rope_dim_list = [16, 56, 56]
+            else:
+                base = max(2, (attention_head_dim // 3) // 2 * 2)
+                rope_dim_list = [base, base, attention_head_dim - base * 2]
+        self.rope_dim_list = [int(dim) for dim in rope_dim_list]
+        if sum(self.rope_dim_list) != attention_head_dim:
+            raise ValueError("sum(rope_dim_list) must equal hidden_size // num_attention_heads.")
+
+        self.img_in = nn.Conv3d(
+            in_channels,
+            hidden_size,
+            kernel_size=self.patch_size,
+            stride=self.patch_size,
+        )
+        self.condition_embedder = JoyImageTimeTextImageEmbedding(
+            dim=hidden_size,
+            time_freq_dim=256,
+            time_proj_dim=hidden_size * 6,
+            text_embed_dim=text_dim,
+        )
+        block_mlp_ratio = mlp_ratio if mlp_ratio is not None else mlp_width_ratio
+        self.double_blocks = nn.ModuleList(
+            [
+                JoyImageTransformerBlock(
+                    dim=hidden_size,
+                    num_attention_heads=num_attention_heads,
+                    attention_head_dim=attention_head_dim,
+                    mlp_width_ratio=block_mlp_ratio,
+                    prefix=f"double_blocks.{i}",
+                )
+                for i in range(num_layers)
+            ]
+        )
+        self.norm_out = FP32LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.proj_out = nn.Linear(hidden_size, self.out_channels * math.prod(self.patch_size))
+        self.gradient_checkpointing = False
+
+    @classmethod
+    def from_config_file(
+        cls,
+        config_path: str | Path,
+        *,
+        od_config: OmniDiffusionConfig | None = None,
+        **overrides: Any,
+    ) -> JoyImageEditTransformer3DModel:
+        with open(config_path) as config_file:
+            config = json.load(config_file)
+        config.update(overrides)
+        return cls(od_config=od_config, **config)
+
+    def get_rotary_pos_embed(
+        self,
+        vis_rope_size: list[int],
+        txt_rope_size: int | None = None,
+    ) -> tuple[
+        tuple[torch.Tensor, torch.Tensor],
+        tuple[torch.Tensor, torch.Tensor] | None,
+    ]:
+        target_ndim = 3
+        if len(vis_rope_size) != target_ndim:
+            vis_rope_size = [1] * (target_ndim - len(vis_rope_size)) + list(vis_rope_size)
+
+        grid = torch.stack(
+            torch.meshgrid(
+                *[torch.linspace(0, size, size + 1, dtype=torch.float32)[:size] for size in vis_rope_size],
+                indexing="ij",
+            ),
+            dim=0,
+        )
+
+        vis_cos: list[torch.Tensor] = []
+        vis_sin: list[torch.Tensor] = []
+        for axis, dim in enumerate(self.rope_dim_list):
+            pos = grid[axis].reshape(-1)
+            freqs = 1.0 / (self.theta ** (torch.arange(0, dim, 2, dtype=torch.float32)[: dim // 2] / dim))
+            freqs = torch.outer(pos.float(), freqs)
+            vis_cos.append(freqs.cos().repeat_interleave(2, dim=1))
+            vis_sin.append(freqs.sin().repeat_interleave(2, dim=1))
+        vis_freqs = (torch.cat(vis_cos, dim=1), torch.cat(vis_sin, dim=1))
+
+        if txt_rope_size is None:
+            return vis_freqs, None
+
+        text_positions = torch.arange(txt_rope_size) + grid.reshape(-1).max().item() + 1
+        txt_cos: list[torch.Tensor] = []
+        txt_sin: list[torch.Tensor] = []
+        for dim in self.rope_dim_list:
+            freqs = 1.0 / (self.theta ** (torch.arange(0, dim, 2, dtype=torch.float32)[: dim // 2] / dim))
+            freqs = torch.outer(text_positions.float(), freqs)
+            txt_cos.append(freqs.cos().repeat_interleave(2, dim=1))
+            txt_sin.append(freqs.sin().repeat_interleave(2, dim=1))
+        txt_freqs = (torch.cat(txt_cos, dim=1), torch.cat(txt_sin, dim=1))
+        return vis_freqs, txt_freqs
+
+    def unpatchify(self, image_tokens: torch.Tensor, t: int, h: int, w: int) -> torch.Tensor:
+        channels = self.out_channels
+        patch_t, patch_h, patch_w = self.patch_size
+        if t * h * w != image_tokens.shape[1]:
+            raise ValueError(f"Expected t*h*w ({t * h * w}) to equal token count ({image_tokens.shape[1]}).")
+        image_tokens = image_tokens.reshape(image_tokens.shape[0], t, h, w, patch_t, patch_h, patch_w, channels)
+        image_tokens = image_tokens.permute(0, 7, 1, 4, 2, 5, 3, 6)
+        return image_tokens.reshape(
+            image_tokens.shape[0],
+            channels,
+            t * patch_t,
+            h * patch_h,
+            w * patch_w,
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        timestep: torch.Tensor | float,
+        encoder_hidden_states: torch.Tensor,
+        encoder_hidden_states_mask: torch.Tensor | None = None,
+        return_dict: bool = True,
+        **_: Any,
+    ) -> tuple[torch.Tensor] | dict[str, torch.Tensor]:
+        is_multi_item = hidden_states.ndim == 6
+        num_items = 0
+        frames_per_item = hidden_states.shape[3] if is_multi_item else None
+        if is_multi_item:
+            num_items = hidden_states.shape[1]
+            if num_items > 1:
+                if self.patch_size[0] != 1:
+                    raise ValueError("For multi-item input, patch_size[0] must be 1.")
+                hidden_states = torch.cat([hidden_states[:, -1:], hidden_states[:, :-1]], dim=1)
+            batch_size, num_items, channels, frames, height, width = hidden_states.shape
+            hidden_states = hidden_states.permute(0, 2, 1, 3, 4, 5).reshape(
+                batch_size,
+                channels,
+                num_items * frames,
+                height,
+                width,
+            )
+        elif hidden_states.ndim != 5:
+            raise ValueError(
+                "JoyImageEditTransformer3DModel expects (B, C, T, H, W) or "
+                f"(B, N, C, T, H, W), got {tuple(hidden_states.shape)}."
+            )
+
+        batch_size, _, original_frames, original_height, original_width = hidden_states.shape
+        patch_t, patch_h, patch_w = self.patch_size
+        if original_frames % patch_t or original_height % patch_h or original_width % patch_w:
+            raise ValueError(
+                "Latent shape must be divisible by patch_size; "
+                f"got {(original_frames, original_height, original_width)} and {self.patch_size}."
+            )
+        patch_frames = original_frames // patch_t
+        patch_height = original_height // patch_h
+        patch_width = original_width // patch_w
+
+        image_tokens = self.img_in(hidden_states).flatten(2).transpose(1, 2)
+
+        if not torch.is_tensor(timestep):
+            timestep = torch.tensor([timestep], device=hidden_states.device, dtype=hidden_states.dtype)
+        timestep = timestep.to(device=hidden_states.device).reshape(-1)
+        if timestep.numel() == 1:
+            timestep = timestep.expand(batch_size)
+
+        _, vector_embedding, text_tokens = self.condition_embedder(
+            timestep,
+            encoder_hidden_states.to(device=hidden_states.device),
+        )
+        if vector_embedding.shape[-1] > self.hidden_size:
+            vector_embedding = vector_embedding.unflatten(1, (6, -1))
+
+        vis_freqs, txt_freqs = self.get_rotary_pos_embed(
+            vis_rope_size=[patch_frames, patch_height, patch_width],
+            txt_rope_size=text_tokens.shape[1] if self.rope_type == "mrope" else None,
+        )
+        vis_freqs = tuple(item.to(hidden_states.device) for item in vis_freqs)
+        if txt_freqs is not None:
+            txt_freqs = tuple(item.to(hidden_states.device) for item in txt_freqs)
+
+        attention_mask = None
+        if encoder_hidden_states_mask is not None:
+            encoder_hidden_states_mask = encoder_hidden_states_mask.to(device=hidden_states.device, dtype=torch.bool)
+            image_mask = torch.ones(
+                batch_size,
+                image_tokens.shape[1],
+                device=hidden_states.device,
+                dtype=torch.bool,
+            )
+            attention_mask = torch.cat([image_mask, encoder_hidden_states_mask], dim=1)
+
+        for block in self.double_blocks:
+            image_tokens, text_tokens = block(
+                hidden_states=image_tokens,
+                encoder_hidden_states=text_tokens,
+                temb=vector_embedding,
+                image_rotary_emb=(vis_freqs, txt_freqs),
+                attention_mask=attention_mask,
+            )
+
+        image_tokens = self.proj_out(self.norm_out(image_tokens))
+        output = self.unpatchify(image_tokens, patch_frames, patch_height, patch_width)
+
+        if is_multi_item:
+            channels = output.shape[1]
+            output = output.reshape(
+                batch_size,
+                channels,
+                num_items,
+                frames_per_item,
+                original_height,
+                original_width,
+            )
+            output = output.permute(0, 2, 1, 3, 4, 5)
+            if num_items > 1:
+                output = torch.cat([output[:, 1:], output[:, :1]], dim=1)
+
+        if return_dict:
+            return {"sample": output}
+        return (output,)
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        params = dict(self.named_parameters())
+        loaded: set[str] = set()
+        for name, loaded_weight in weights:
+            mapped_name = self._map_weight_name(name)
+            param = params.get(mapped_name)
+            if param is None:
+                logger.warning(
+                    "Skipping JoyAI-Image-Edit transformer weight %s mapped to "
+                    "%s -- not found in model parameters. This may indicate an "
+                    "incomplete implementation or checkpoint mismatch.",
+                    name,
+                    mapped_name,
+                )
+                continue
+            default_weight_loader(param, loaded_weight)
+            loaded.add(mapped_name)
+        return loaded
+
+    @staticmethod
+    def _map_weight_name(name: str) -> str:
+        if name.startswith("transformer."):
+            name = name[len("transformer.") :]
+        prefix_replacements = {
+            "patch_embedding.": "img_in.",
+            "time_embed.": "condition_embedder.time_embedder.",
+            "time_text_embed.": "condition_embedder.",
+            "text_embedder.": "condition_embedder.text_embedder.",
+            "final_layer.linear.": "proj_out.",
+            "final_layer.norm_final.": "norm_out.",
+        }
+        for source, target in prefix_replacements.items():
+            if name.startswith(source):
+                return target + name[len(source) :]
+        return name

vllm_omni/diffusion/models/joy_image/pipeline_joy_image_edit.py

@@ -0,0 +1,953 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+# Adapted from Hugging Face Diffusers commit
+# 23ba73e1d2079c4b89959484ed0ca1c22e7ef998:
+# src/diffusers/pipelines/joyimage/pipeline_joyimage_edit.py
+
+from __future__ import annotations
+
+import inspect
+import math
+import os
+from collections.abc import Callable
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, cast
+
+import numpy as np
+import PIL.Image
+import torch
+from diffusers.image_processor import VaeImageProcessor
+from diffusers.schedulers.scheduling_flow_match_euler_discrete import (
+    FlowMatchEulerDiscreteScheduler,
+)
+from diffusers.utils.torch_utils import randn_tensor
+from torch import nn
+from transformers import Qwen2TokenizerFast, Qwen3VLForConditionalGeneration, Qwen3VLProcessor
+from vllm.model_executor.models.utils import AutoWeightsLoader
+
+from vllm_omni.diffusion.data import DiffusionOutput
+from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import (
+    DistributedAutoencoderKLWan,
+)
+from vllm_omni.diffusion.distributed.utils import get_local_device
+from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+from vllm_omni.diffusion.model_loader.hub_prefetch import (
+    from_pretrained_with_prefetch,
+    prefetch_subfolders,
+)
+from vllm_omni.diffusion.models.interface import (
+    SupportImageInput,
+    SupportsComponentDiscovery,
+)
+from vllm_omni.diffusion.models.joy_image.cfg_parallel import (
+    JoyImageEditCFGParallelMixin,
+)
+from vllm_omni.diffusion.models.joy_image.joy_image_edit_transformer import (
+    JoyImageEditTransformer3DModel,
+    _raise_if_unsupported_tensor_parallel,
+)
+from vllm_omni.diffusion.offloader.config import (
+    OffloadStrategy,
+    offload_enabled,
+    resolve_offload_strategy,
+)
+from vllm_omni.diffusion.offloader.sequential_backend import SequentialOffloadHook
+from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import (
+    DiffusionPipelineProfilerMixin,
+)
+from vllm_omni.model_executor.model_loader.weight_utils import (
+    download_weights_from_hf_specific,
+)
+from vllm_omni.platforms import current_omni_platform
+
+if TYPE_CHECKING:
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+    from vllm_omni.diffusion.request import OmniDiffusionRequest
+
+
+JOY_MAX_IMAGE_SEQ_LEN = 4096
+# Reference: Hugging Face Diffusers commit 23ba73e1d2079c4b89959484ed0ca1c22e7ef998,
+# src/diffusers/pipelines/joyimage/pipeline_joyimage_edit.py.
+JOY_BUCKETS = [
+    (512, 1792),
+    (512, 1856),
+    (512, 1920),
+    (512, 1984),
+    (512, 2048),
+    (576, 1600),
+    (576, 1664),
+    (576, 1728),
+    (576, 1792),
+    (640, 1472),
+    (640, 1536),
+    (640, 1600),
+    (704, 1344),
+    (704, 1408),
+    (704, 1472),
+    (768, 1216),
+    (768, 1280),
+    (768, 1344),
+    (832, 1152),
+    (832, 1216),
+    (896, 1088),
+    (896, 1152),
+    (960, 1024),
+    (960, 1088),
+    (1024, 960),
+    (1024, 1024),
+    (1088, 896),
+    (1088, 960),
+    (1152, 832),
+    (1152, 896),
+    (1216, 768),
+    (1216, 832),
+    (1280, 768),
+    (1344, 704),
+    (1344, 768),
+    (1408, 704),
+    (1472, 640),
+    (1472, 704),
+    (1536, 640),
+    (1600, 576),
+    (1600, 640),
+    (1664, 576),
+    (1728, 576),
+    (1792, 512),
+    (1792, 576),
+    (1856, 512),
+    (1920, 512),
+    (1984, 512),
+    (2048, 512),
+]
+JOY_PROMPT_TEMPLATE_START_IDX = 34
+# The upstream template is named "multiple_images", but this port currently
+# inserts a single reference-image token and rejects multi-image prompts.
+JOY_PROMPT_TEMPLATE = (
+    "<|im_start|>system\n \\nDescribe the image by detailing the color, shape, size, texture, "
+    "quantity, text, spatial relationships of the objects and background:<|im_end|>\n"
+    "{}<|im_start|>assistant\n"
+)
+JOY_VISION_TOKEN = "<|vision_start|><|image_pad|><|vision_end|>"
+
+
+def _get_model_path(model_name: str) -> str:
+    if os.path.exists(model_name):
+        return model_name
+    return download_weights_from_hf_specific(
+        model_name,
+        None,
+        ["vae/config.json", "transformer/config.json"],
+        require_all=True,
+    )
+
+
+def _get_transformer_config_kwargs_from_od_config(
+    od_config: OmniDiffusionConfig,
+) -> dict[str, Any]:
+    tf_model_config = getattr(od_config, "tf_model_config", None)
+    to_dict = getattr(tf_model_config, "to_dict", None)
+    if not callable(to_dict):
+        return {}
+    return {key: value for key, value in to_dict().items() if not key.startswith("_")}
+
+
+def _should_defer_component_device_placement(
+    od_config: OmniDiffusionConfig,
+) -> bool:
+    return offload_enabled(od_config)
+
+
+def _uses_model_level_cpu_offload(
+    od_config: OmniDiffusionConfig,
+) -> bool:
+    return resolve_offload_strategy(od_config) is OffloadStrategy.MODEL_LEVEL
+
+
+def _raise_if_unsupported_hsdp(od_config: OmniDiffusionConfig) -> None:
+    parallel_config = getattr(od_config, "parallel_config", None)
+    if getattr(parallel_config, "use_hsdp", False):
+        raise ValueError(
+            "JoyImageEditPipeline does not support HSDP yet. "
+            "Please disable `use_hsdp`, or add `_hsdp_shard_conditions` and HSDP parity tests."
+        )
+
+
+def _format_qwen_multimodal_prompt(prompt: str) -> str:
+    prompt = f"<image>\n{prompt}"
+    prompt = f"<|im_start|>user\n{prompt}<|im_end|>\n"
+    return prompt.replace("<image>\n", JOY_VISION_TOKEN)
+
+
+def _find_best_bucket(height: int, width: int) -> tuple[int, int]:
+    target_ratio = height / width
+    return min(JOY_BUCKETS, key=lambda size: abs(size[0] / size[1] - target_ratio))
+
+
+def _resize_center_crop(
+    image: PIL.Image.Image,
+    height: int,
+    width: int,
+) -> PIL.Image.Image:
+    image = image.convert("RGB")
+    source_width, source_height = image.size
+    scale = max(height / source_height, width / source_width)
+    resized_height = math.ceil(source_height * scale)
+    resized_width = math.ceil(source_width * scale)
+    image = image.resize((resized_width, resized_height), PIL.Image.Resampling.BILINEAR)
+    left = (resized_width - width) // 2
+    top = (resized_height - height) // 2
+    return image.crop((left, top, left + width, top + height))
+
+
+def _pil_to_tensor(image: PIL.Image.Image) -> torch.Tensor:
+    array = np.asarray(image, dtype=np.float32) / 127.5 - 1.0
+    tensor = torch.from_numpy(array).permute(2, 0, 1).unsqueeze(0)
+    return tensor.unsqueeze(2).contiguous()
+
+
+def _extract_single_image(
+    prompt: dict[str, Any] | str,
+) -> PIL.Image.Image | torch.Tensor | np.ndarray:
+    multi_modal_data = prompt.get("multi_modal_data", {}) if not isinstance(prompt, str) else {}
+    raw_image = multi_modal_data.get("image")
+    if raw_image is None or (isinstance(raw_image, list) and len(raw_image) == 0):
+        raise ValueError("Received no input image. JoyAI-Image-Edit requires exactly one input image.")
+    if isinstance(raw_image, list):
+        if len(raw_image) != 1:
+            raise ValueError("Received multiple input images. JoyAI-Image-Edit v1 supports exactly one image.")
+        raw_image = raw_image[0]
+    if isinstance(raw_image, str):
+        return PIL.Image.open(raw_image)
+    return cast(PIL.Image.Image | torch.Tensor | np.ndarray, raw_image)
+
+
+def _image_size(image: PIL.Image.Image | torch.Tensor | np.ndarray) -> tuple[int, int]:
+    if isinstance(image, PIL.Image.Image):
+        return image.size
+    if isinstance(image, torch.Tensor):
+        height, width = image.shape[-2:]
+        return int(width), int(height)
+    height, width = image.shape[:2]
+    return int(width), int(height)
+
+
+def _to_pil_image(
+    image: PIL.Image.Image | torch.Tensor | np.ndarray,
+) -> PIL.Image.Image:
+    if isinstance(image, PIL.Image.Image):
+        return image
+    if isinstance(image, torch.Tensor):
+        tensor = image.detach().cpu()
+        if tensor.ndim == 4:
+            tensor = tensor[0]
+        if tensor.ndim == 3 and tensor.shape[0] in (1, 3):
+            tensor = tensor.permute(1, 2, 0)
+        array = tensor.float().numpy()
+    else:
+        array = image
+    if array.dtype != np.uint8:
+        if array.min() < 0:
+            array = (array + 1.0) / 2.0
+        array = np.clip(array * 255.0 if array.max() <= 1.0 else array, 0, 255).astype(np.uint8)
+    return PIL.Image.fromarray(array).convert("RGB")
+
+
+def _cast_floating_model_inputs(model_inputs: Any, dtype: torch.dtype) -> Any:
+    for key in ("pixel_values", "pixel_values_videos"):
+        get_value = getattr(model_inputs, "get", None)
+        tensor = get_value(key) if callable(get_value) else getattr(model_inputs, key, None)
+        if isinstance(tensor, torch.Tensor) and torch.is_floating_point(tensor):
+            model_inputs[key] = tensor.to(dtype=dtype)
+    return model_inputs
+
+
+def _get_single_request_prompts(request: OmniDiffusionRequest) -> tuple[list[Any], bool]:
+    prompts = getattr(request, "prompts", None)
+    if prompts is not None:
+        return list(prompts), True
+    return [request.prompt], False
+
+
+def _set_single_request_prompt(
+    request: OmniDiffusionRequest,
+    prompts: list[Any],
+    *,
+    uses_prompts_attr: bool,
+) -> None:
+    if uses_prompts_attr:
+        request.prompts = prompts
+    else:
+        request.prompt = prompts[0]
+
+
+# Normalize the single reference image into a supported Joy bucket before
+# forward(); default and explicitly requested sizes are both snapped to the
+# nearest bucket, so the final output size can differ from user-provided
+# height/width while keeping VAE/DiT latent grids within trained resolutions.
+def get_joy_image_edit_pre_process_func(
+    od_config: OmniDiffusionConfig,
+) -> Callable[[OmniDiffusionRequest], OmniDiffusionRequest]:
+    def pre_process_func(request: OmniDiffusionRequest) -> OmniDiffusionRequest:
+        prompts, uses_prompts_attr = _get_single_request_prompts(request)
+        if len(prompts) != 1:
+            raise ValueError("JoyAI-Image-Edit v1 supports exactly one prompt per request.")
+        for prompt_index, prompt in enumerate(prompts):
+            if isinstance(prompt, str):
+                prompt = {"prompt": prompt}
+            prompt.setdefault("additional_information", {})
+            image = _extract_single_image(prompt)
+            original_width, original_height = _image_size(image)
+            height_is_set = request.sampling_params.height is not None
+            width_is_set = request.sampling_params.width is not None
+            if height_is_set != width_is_set:
+                raise ValueError("JoyAI-Image-Edit requires both `height` and `width`, or neither.")
+            if not height_is_set and not width_is_set:
+                height, width = _find_best_bucket(original_height, original_width)
+            else:
+                height, width = _find_best_bucket(
+                    request.sampling_params.height,
+                    request.sampling_params.width,
+                )
+
+            prompt_image = _resize_center_crop(_to_pil_image(image), height, width)
+            image_tensor = _pil_to_tensor(prompt_image)
+            prompt["additional_information"].update(
+                {
+                    "image_tensor": image_tensor,
+                    "prompt_image": prompt_image,
+                    "height": height,
+                    "width": width,
+                    "original_size": (original_width, original_height),
+                    "resized_size": (width, height),
+                }
+            )
+            request.sampling_params.height = height
+            request.sampling_params.width = width
+            prompts[prompt_index] = prompt
+        _set_single_request_prompt(request, prompts, uses_prompts_attr=uses_prompts_attr)
+        return request
+
+    return pre_process_func
+
+
+def get_joy_image_edit_post_process_func(
+    od_config: OmniDiffusionConfig,
+) -> Callable[[torch.Tensor], list[PIL.Image.Image]]:
+    image_processor = VaeImageProcessor()
+
+    def post_process_func(images: torch.Tensor) -> list[PIL.Image.Image]:
+        return image_processor.postprocess(images, output_type="pil")
+
+    return post_process_func
+
+
+def retrieve_latents(encoder_output: Any, generator: torch.Generator | None = None) -> torch.Tensor:
+    if hasattr(encoder_output, "latent_dist"):
+        return encoder_output.latent_dist.sample(generator=generator)
+    if hasattr(encoder_output, "latents"):
+        return encoder_output.latents
+    raise AttributeError("Could not access latents of provided VAE encoder output.")
+
+
+def retrieve_timesteps(
+    scheduler: Any,
+    num_inference_steps: int | None = None,
+    device: str | torch.device | None = None,
+    timesteps: list[int] | torch.Tensor | None = None,
+    sigmas: list[float] | None = None,
+) -> tuple[torch.Tensor, int]:
+    if timesteps is not None and sigmas is not None:
+        raise ValueError("Only one of `timesteps` or `sigmas` can be provided.")
+
+    parameters = inspect.signature(scheduler.set_timesteps).parameters
+    if timesteps is not None:
+        if "timesteps" not in parameters:
+            raise ValueError(f"{scheduler.__class__.__name__} does not support custom `timesteps`.")
+        scheduler.set_timesteps(timesteps=timesteps, device=device)
+    elif sigmas is not None:
+        if "sigmas" not in parameters:
+            raise ValueError(f"{scheduler.__class__.__name__} does not support custom `sigmas`.")
+        scheduler.set_timesteps(sigmas=sigmas, device=device)
+    else:
+        scheduler.set_timesteps(num_inference_steps, device=device)
+
+    resolved_timesteps = scheduler.timesteps
+    return resolved_timesteps, len(resolved_timesteps)
+
+
+class JoyImageEditPipeline(
+    nn.Module,
+    SupportImageInput,
+    JoyImageEditCFGParallelMixin,
+    DiffusionPipelineProfilerMixin,
+    SupportsComponentDiscovery,
+):
+    _dit_modules = ["transformer"]
+    # Joy's VAE is large enough that model-level CPU offload must treat it as
+    # mutually exclusive with the DiT, instead of pinning it on GPU.
+    _encoder_modules = ["text_encoder", "vae"]
+    _vae_modules: list[str] = []
+    _resident_modules: list[str] = []
+
+    def __init__(
+        self,
+        *,
+        od_config: OmniDiffusionConfig,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        # Joy has no `_hsdp_shard_conditions` yet, so fail before generic
+        # HSDP setup reaches the transformer.
+        _raise_if_unsupported_hsdp(od_config)
+        _raise_if_unsupported_tensor_parallel(od_config)
+        self.od_config = od_config
+        self.weights_sources = [
+            DiffusersPipelineLoader.ComponentSource(
+                model_or_path=od_config.model,
+                subfolder="transformer",
+                revision=None,
+                prefix="transformer.",
+                fall_back_to_pt=True,
+            )
+        ]
+        self.device = get_local_device()
+        model = od_config.model
+        local_files_only = os.path.exists(model)
+        joy_subfolders = [
+            "scheduler",
+            "text_encoder",
+            "vae",
+            "tokenizer",
+            "processor",
+            "transformer",
+        ]
+        prefetch_subfolders(
+            model,
+            joy_subfolders,
+            local_files_only=local_files_only,
+        )
+        self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
+            model,
+            subfolder="scheduler",
+            local_files_only=local_files_only,
+        )
+        self.processor = from_pretrained_with_prefetch(
+            Qwen3VLProcessor.from_pretrained,
+            model,
+            subfolder="processor",
+            prefetch_list=joy_subfolders,
+            local_files_only=local_files_only,
+        )
+        self.tokenizer = getattr(self.processor, "tokenizer", None)
+        if self.tokenizer is None:
+            self.tokenizer = Qwen2TokenizerFast.from_pretrained(
+                model,
+                subfolder="tokenizer",
+                local_files_only=local_files_only,
+            )
+        defer_component_device_placement = _should_defer_component_device_placement(od_config)
+        self.text_encoder = from_pretrained_with_prefetch(
+            Qwen3VLForConditionalGeneration.from_pretrained,
+            model,
+            subfolder="text_encoder",
+            prefetch_list=joy_subfolders,
+            torch_dtype=torch.bfloat16,
+            local_files_only=local_files_only,
+        )
+        if not defer_component_device_placement:
+            self.text_encoder = self.text_encoder.to(self.device)
+        self.vae = from_pretrained_with_prefetch(
+            DistributedAutoencoderKLWan.from_pretrained,
+            model,
+            subfolder="vae",
+            prefetch_list=joy_subfolders,
+            torch_dtype=torch.bfloat16,
+            local_files_only=local_files_only,
+        )
+        if not defer_component_device_placement:
+            self.vae = self.vae.to(self.device)
+        transformer_config_kwargs = _get_transformer_config_kwargs_from_od_config(od_config)
+        if transformer_config_kwargs:
+            self.transformer = JoyImageEditTransformer3DModel(
+                od_config=od_config,
+                **transformer_config_kwargs,
+            )
+        else:
+            model_path = _get_model_path(model)
+            transformer_config_path = Path(model_path) / "transformer" / "config.json"
+            if not transformer_config_path.exists():
+                raise FileNotFoundError(
+                    "JoyAI-Image-Edit requires transformer/config.json to "
+                    f"instantiate the DiT; could not find {transformer_config_path}."
+                )
+            self.transformer = JoyImageEditTransformer3DModel.from_config_file(
+                transformer_config_path,
+                od_config=od_config,
+            )
+
+        self.vae_scale_factor = int(getattr(self.vae.config, "scale_factor_spatial", 8))
+        self.latent_channels = int(getattr(self.vae.config, "z_dim", 16))
+        self.tokenizer_max_length = JOY_MAX_IMAGE_SEQ_LEN
+        self.prompt_template_encode = JOY_PROMPT_TEMPLATE
+        self.prompt_template_encode_start_idx = JOY_PROMPT_TEMPLATE_START_IDX
+        self.setup_diffusion_pipeline_profiler(
+            enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler
+        )
+
+    @staticmethod
+    def _extract_masked_hidden(hidden_states: torch.Tensor, mask: torch.Tensor) -> list[torch.Tensor]:
+        bool_mask = mask.bool()
+        valid_lengths = bool_mask.sum(dim=1)
+        selected = hidden_states[bool_mask]
+        return list(torch.split(selected, valid_lengths.tolist(), dim=0))
+
+    @staticmethod
+    def _pad_prompt_embeds(
+        split_hidden_states: list[torch.Tensor],
+        max_sequence_length: int,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if max_sequence_length <= 0:
+            raise ValueError("max_sequence_length must be greater than 0.")
+        split_hidden_states = [item[-max_sequence_length:] for item in split_hidden_states]
+        max_seq_len = max(item.shape[0] for item in split_hidden_states)
+        prompt_embeds = torch.stack(
+            [
+                torch.cat(
+                    [item, item.new_zeros(max_seq_len - item.shape[0], item.shape[1])],
+                    dim=0,
+                )
+                for item in split_hidden_states
+            ],
+            dim=0,
+        )
+        prompt_mask = torch.stack(
+            [
+                torch.cat(
+                    [
+                        torch.ones(item.shape[0], dtype=torch.long, device=item.device),
+                        torch.zeros(
+                            max_seq_len - item.shape[0],
+                            dtype=torch.long,
+                            device=item.device,
+                        ),
+                    ]
+                )
+                for item in split_hidden_states
+            ],
+            dim=0,
+        )
+        return prompt_embeds, prompt_mask
+
+    def _get_last_layer_pre_norm_hidden(
+        self,
+        model_inputs: Any,
+    ) -> torch.Tensor:
+        captured: dict[str, torch.Tensor] = {}
+
+        def hook(_module: nn.Module, _args: tuple[Any, ...], output: Any) -> None:
+            captured["hidden_states"] = output[0] if isinstance(output, tuple) else output
+
+        handle = self.text_encoder.model.language_model.layers[-1].register_forward_hook(hook)
+        try:
+            input_items = getattr(model_inputs, "items", None)
+            if callable(input_items):
+                text_encoder_inputs = dict(input_items())
+            else:
+                text_encoder_inputs = dict(vars(model_inputs))
+            text_encoder_inputs["output_hidden_states"] = False
+            self.text_encoder(**text_encoder_inputs)
+        finally:
+            handle.remove()
+        if "hidden_states" not in captured:
+            raise RuntimeError("Failed to capture Qwen3VL last-layer pre-norm hidden states.")
+        return captured["hidden_states"]
+
+    def _get_qwen_prompt_embeds(
+        self,
+        prompt: str | list[str],
+        image: PIL.Image.Image | list[PIL.Image.Image],
+        dtype: torch.dtype | None = None,
+        max_sequence_length: int | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        dtype = dtype or self.text_encoder.dtype
+        prompt_list = [prompt] if isinstance(prompt, str) else prompt
+        image_list = [image] if isinstance(image, PIL.Image.Image) else image
+        texts = [_format_qwen_multimodal_prompt(item) for item in prompt_list]
+        texts = [self.prompt_template_encode.format(item) for item in texts]
+        model_inputs = self.processor(
+            text=texts,
+            images=image_list,
+            padding=True,
+            return_tensors="pt",
+        ).to(self.device)
+        model_inputs = _cast_floating_model_inputs(model_inputs, dtype)
+        hidden_states = self._get_last_layer_pre_norm_hidden(model_inputs)
+        split_hidden_states = self._extract_masked_hidden(hidden_states, model_inputs.attention_mask)
+        split_hidden_states = [item[self.prompt_template_encode_start_idx :] for item in split_hidden_states]
+        prompt_embeds, prompt_mask = self._pad_prompt_embeds(
+            split_hidden_states,
+            self.tokenizer_max_length if max_sequence_length is None else max_sequence_length,
+        )
+        return prompt_embeds.to(dtype=dtype, device=self.device), prompt_mask.to(device=self.device)
+
+    @staticmethod
+    def _validate_prompt_embeds_and_mask(
+        prompt_embeds: torch.Tensor,
+        prompt_embeds_mask: torch.Tensor,
+        *,
+        embeds_name: str = "prompt_embeds",
+        mask_name: str = "prompt_embeds_mask",
+    ) -> None:
+        """Fail early for precomputed prompt embeds before attention uses them."""
+        if prompt_embeds.ndim != 3:
+            raise ValueError(
+                f"{embeds_name} must be a 3D tensor of shape "
+                f"(batch, sequence_length, hidden_size), got shape {tuple(prompt_embeds.shape)}."
+            )
+        if prompt_embeds_mask.ndim != 2:
+            raise ValueError(
+                f"{mask_name} must be a 2D tensor of shape "
+                f"(batch, sequence_length), got shape {tuple(prompt_embeds_mask.shape)}."
+            )
+        if prompt_embeds.shape[0] != prompt_embeds_mask.shape[0]:
+            raise ValueError(
+                f"{embeds_name} and {mask_name} must have the same batch size, got "
+                f"{prompt_embeds.shape[0]} and {prompt_embeds_mask.shape[0]}."
+            )
+        if prompt_embeds.shape[1] != prompt_embeds_mask.shape[1]:
+            raise ValueError(
+                f"{embeds_name} and {mask_name} must have the same sequence length, got "
+                f"{prompt_embeds.shape[1]} and {prompt_embeds_mask.shape[1]}."
+            )
+
+    def encode_prompt(
+        self,
+        prompt: str | list[str],
+        image: PIL.Image.Image | list[PIL.Image.Image],
+        num_images_per_prompt: int,
+        prompt_embeds: torch.Tensor | None = None,
+        prompt_embeds_mask: torch.Tensor | None = None,
+        max_sequence_length: int | None = None,
+        embeds_name: str = "prompt_embeds",
+        mask_name: str = "prompt_embeds_mask",
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        prompt_list = [prompt] if isinstance(prompt, str) else prompt
+        if prompt_embeds is None:
+            prompt_embeds, prompt_embeds_mask = self._get_qwen_prompt_embeds(
+                prompt_list,
+                image,
+                max_sequence_length=max_sequence_length,
+            )
+        if prompt_embeds_mask is None:
+            raise ValueError(f"{mask_name} must be provided with {embeds_name}.")
+        self._validate_prompt_embeds_and_mask(
+            prompt_embeds,
+            prompt_embeds_mask,
+            embeds_name=embeds_name,
+            mask_name=mask_name,
+        )
+        if max_sequence_length is not None:
+            if max_sequence_length <= 0:
+                raise ValueError("max_sequence_length must be greater than 0.")
+            prompt_embeds = prompt_embeds[:, -max_sequence_length:]
+            prompt_embeds_mask = prompt_embeds_mask[:, -max_sequence_length:]
+        prompt_embeds = prompt_embeds.repeat_interleave(num_images_per_prompt, dim=0)
+        prompt_embeds_mask = prompt_embeds_mask.repeat_interleave(num_images_per_prompt, dim=0)
+        return prompt_embeds, prompt_embeds_mask
+
+    def _latent_stats(self, device: torch.device, dtype: torch.dtype) -> tuple[torch.Tensor, torch.Tensor]:
+        latents_mean = (
+            torch.tensor(self.vae.config.latents_mean)
+            .view(1, self.latent_channels, 1, 1, 1)
+            .to(
+                device=device,
+                dtype=dtype,
+            )
+        )
+        latents_std = (
+            torch.tensor(self.vae.config.latents_std)
+            .view(1, self.latent_channels, 1, 1, 1)
+            .to(
+                device=device,
+                dtype=dtype,
+            )
+        )
+        return latents_mean, latents_std
+
+    def _encode_vae_image(
+        self,
+        image: torch.Tensor,
+        generator: torch.Generator | list[torch.Generator] | None,
+    ):
+        if isinstance(generator, list):
+            image_latents = [
+                retrieve_latents(self.vae.encode(image[item_index : item_index + 1]))
+                for item_index in range(image.shape[0])
+            ]
+            image_latents = torch.cat(image_latents, dim=0)
+        else:
+            image_latents = retrieve_latents(self.vae.encode(image))
+        latents_mean, latents_std = self._latent_stats(image_latents.device, image_latents.dtype)
+        return (image_latents - latents_mean) / latents_std
+
+    def _prepare_latents(
+        self,
+        image: torch.Tensor,
+        batch_size: int,
+        num_channels_latents: int,
+        height: int,
+        width: int,
+        dtype: torch.dtype,
+        device: torch.device,
+        generator: torch.Generator | list[torch.Generator] | None,
+        latents: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        latent_height = int(height) // self.vae_scale_factor
+        latent_width = int(width) // self.vae_scale_factor
+        expected_noise_shape = (
+            batch_size,
+            1,
+            num_channels_latents,
+            1,
+            latent_height,
+            latent_width,
+        )
+        if isinstance(generator, list) and len(generator) != batch_size:
+            raise ValueError(f"Generator list length {len(generator)} must match effective batch size {batch_size}.")
+
+        self._prepare_vae_for_decode()
+        image = image.to(device=device, dtype=dtype)
+        image_latents = self._encode_vae_image(image, generator)
+        image_latents = image_latents.to(device=device, dtype=dtype)
+        expected_image_tail = (
+            num_channels_latents,
+            1,
+            latent_height,
+            latent_width,
+        )
+        if image_latents.ndim != 5 or tuple(image_latents.shape[1:]) != expected_image_tail:
+            raise ValueError(
+                "Joy image latents must have shape "
+                f"(B, {expected_image_tail[0]}, {expected_image_tail[1]}, "
+                f"{expected_image_tail[2]}, {expected_image_tail[3]}), "
+                f"got {tuple(image_latents.shape)}."
+            )
+        image_latents = image_latents.unsqueeze(1)
+        if batch_size > image_latents.shape[0]:
+            if batch_size % image_latents.shape[0] != 0:
+                raise ValueError(
+                    f"Cannot duplicate image batch size {image_latents.shape[0]} to effective batch {batch_size}."
+                )
+            repeats = batch_size // image_latents.shape[0]
+            image_latents = image_latents.repeat_interleave(repeats, dim=0)
+        elif batch_size < image_latents.shape[0]:
+            raise ValueError(f"Image batch size {image_latents.shape[0]} exceeds effective batch size {batch_size}.")
+
+        if latents is None:
+            noise_latents = randn_tensor(
+                expected_noise_shape,
+                generator=generator,
+                device=device,
+                dtype=dtype,
+            )
+        else:
+            noise_latents = latents.to(device=device, dtype=dtype)
+            if noise_latents.ndim == 5:
+                noise_latents = noise_latents.unsqueeze(1)
+            if tuple(noise_latents.shape) != expected_noise_shape:
+                raise ValueError(
+                    f"Joy noise latents must have shape {expected_noise_shape}, got {tuple(noise_latents.shape)}."
+                )
+        return torch.cat([image_latents, noise_latents], dim=1), image_latents
+
+    def _offload_transformer_if_model_level(self) -> None:
+        od_config = getattr(self, "od_config", None)
+        if od_config is None or not _uses_model_level_cpu_offload(od_config):
+            return
+
+        transformer = getattr(self, "transformer", None)
+        if isinstance(transformer, nn.Module):
+            pin_memory = getattr(od_config, "pin_cpu_memory", True)
+            SequentialOffloadHook._move_params(
+                transformer,
+                torch.device("cpu"),
+                non_blocking=pin_memory,
+                pin_memory=pin_memory,
+            )
+            current_omni_platform.empty_cache()
+
+    def _prepare_vae_for_decode(self) -> None:
+        od_config = getattr(self, "od_config", None)
+        if od_config is None or not _should_defer_component_device_placement(od_config):
+            return
+
+        self._offload_transformer_if_model_level()
+        SequentialOffloadHook._move_params(self.vae, self.device, non_blocking=False)
+        current_omni_platform.synchronize()
+
+    def _vae_device(self, fallback: torch.device) -> torch.device:
+        try:
+            return next(self.vae.parameters()).device
+        except (AttributeError, StopIteration):
+            return fallback
+
+    def _decode_latents(self, latents: torch.Tensor) -> torch.Tensor:
+        self._prepare_vae_for_decode()
+        batch_size, num_items, channels, frames, height, width = latents.shape
+        flat_latents = latents.reshape(batch_size * num_items, channels, frames, height, width)
+        flat_latents = flat_latents.to(device=self._vae_device(flat_latents.device))
+        latents_mean, latents_std = self._latent_stats(flat_latents.device, flat_latents.dtype)
+        flat_latents = flat_latents * latents_std + latents_mean
+        decoded = self.vae.decode(flat_latents, return_dict=False)[0]
+        decoded = decoded.reshape(
+            batch_size,
+            num_items,
+            frames,
+            decoded.shape[1],
+            decoded.shape[-2],
+            decoded.shape[-1],
+        )
+        return decoded[:, -1, 0]
+
+    @property
+    def interrupt(self) -> bool:
+        return getattr(self, "_interrupt", False)
+
+    def forward(
+        self,
+        req: OmniDiffusionRequest,
+        prompt_embeds: torch.Tensor | None = None,
+        prompt_embeds_mask: torch.Tensor | None = None,
+        negative_prompt_embeds: torch.Tensor | None = None,
+        negative_prompt_embeds_mask: torch.Tensor | None = None,
+        generator: torch.Generator | list[torch.Generator] | None = None,
+        latents: torch.Tensor | None = None,
+        output_type: str | None = "pil",
+    ) -> DiffusionOutput:
+        # Step 1: read the preprocessed single-image edit request.
+        if len(req.prompts) != 1:
+            raise ValueError("JoyImageEditPipeline supports exactly one prompt per request.")
+        first_prompt = req.prompts[0]
+        prompt = first_prompt if isinstance(first_prompt, str) else (first_prompt.get("prompt") or "")
+        additional_information = {} if isinstance(first_prompt, str) else first_prompt.get("additional_information", {})
+        image = additional_information.get("image_tensor")
+        prompt_image = additional_information.get("prompt_image")
+        height = additional_information.get("height") or req.sampling_params.height
+        width = additional_information.get("width") or req.sampling_params.width
+        if image is None or prompt_image is None or height is None or width is None:
+            raise ValueError("JoyImageEditPipeline requires preprocessed image information in the request.")
+
+        num_inference_steps = req.sampling_params.num_inference_steps or 50
+        timesteps = req.sampling_params.timesteps
+        sigmas = req.sampling_params.sigmas
+        tokenizer_max_length = getattr(self, "tokenizer_max_length", JOY_MAX_IMAGE_SEQ_LEN)
+        max_sequence_length = (
+            tokenizer_max_length
+            if req.sampling_params.max_sequence_length is None
+            else req.sampling_params.max_sequence_length
+        )
+        generator = req.sampling_params.generator or generator
+        num_images_per_prompt = max(req.sampling_params.num_outputs_per_prompt, 1)
+        true_cfg_scale = self.resolve_effective_true_cfg_scale(req)
+        latents = req.sampling_params.latents if req.sampling_params.latents is not None else latents
+        output_type = req.sampling_params.output_type or output_type or "pil"
+        if output_type not in {"pil", "latent"}:
+            raise ValueError("JoyImageEditPipeline supports only `pil` and `latent` output types.")
+        negative_prompt = "" if isinstance(first_prompt, str) else first_prompt.get("negative_prompt") or ""
+        do_true_cfg = true_cfg_scale > 1.0
+
+        self._current_timestep = None
+        self._interrupt = False
+
+        # Step 2: encode text plus the reference image through Qwen3-VL
+        # (including its vision/ViT path) to produce DiT conditioning.
+        prompt_embeds, prompt_embeds_mask = self.encode_prompt(
+            prompt,
+            prompt_image,
+            num_images_per_prompt=num_images_per_prompt,
+            prompt_embeds=prompt_embeds,
+            prompt_embeds_mask=prompt_embeds_mask,
+            max_sequence_length=max_sequence_length,
+        )
+        if do_true_cfg:
+            negative_prompt_embeds, negative_prompt_embeds_mask = self.encode_prompt(
+                negative_prompt,
+                prompt_image,
+                num_images_per_prompt=num_images_per_prompt,
+                prompt_embeds=negative_prompt_embeds,
+                prompt_embeds_mask=negative_prompt_embeds_mask,
+                max_sequence_length=max_sequence_length,
+                embeds_name="negative_prompt_embeds",
+                mask_name="negative_prompt_embeds_mask",
+            )
+        else:
+            negative_prompt_embeds = None
+            negative_prompt_embeds_mask = None
+
+        batch_size = prompt_embeds.shape[0]
+        if req.sampling_params.seed is not None and generator is None:
+            generator = torch.Generator(device=self.device).manual_seed(req.sampling_params.seed)
+
+        # Step 3: VAE-encode the reference image and create the target noise
+        # latent; the latent stack is [reference image latents, target noise].
+        latents, image_latents = self._prepare_latents(
+            image=image,
+            batch_size=batch_size,
+            num_channels_latents=self.latent_channels,
+            height=height,
+            width=width,
+            dtype=prompt_embeds.dtype,
+            device=self.device,
+            generator=generator,
+            latents=latents,
+        )
+        # Step 4: build the scheduler timestep/sigma schedule.
+        timesteps, _ = retrieve_timesteps(
+            self.scheduler,
+            num_inference_steps=num_inference_steps,
+            device=self.device,
+            timesteps=timesteps,
+            sigmas=sigmas,
+        )
+        # Step 5: run the DiT denoising loop. The reference image latents stay
+        # fixed while the target noise latent is updated each scheduler step.
+        latents = self.diffuse(
+            latents=latents,
+            image_latents=image_latents,
+            prompt_embeds=prompt_embeds,
+            prompt_embeds_mask=prompt_embeds_mask,
+            negative_prompt_embeds=negative_prompt_embeds,
+            negative_prompt_embeds_mask=negative_prompt_embeds_mask,
+            timesteps=timesteps,
+            do_true_cfg=do_true_cfg,
+            true_cfg_scale=true_cfg_scale,
+            # JoyImage's upstream pipeline always applies CFG norm rescaling.
+            cfg_normalize=True,
+        )
+        if output_type == "latent" or req.is_dummy_run():
+            self._offload_transformer_if_model_level()
+            return DiffusionOutput(output=latents[:, -1].detach().cpu())
+        # Step 6: decode only the target latent back to image space.
+        images = self._decode_latents(latents)
+        return DiffusionOutput(output=images)
+
+    def load_weights(self, weights):
+        return AutoWeightsLoader(self).load_weights(weights)
+
+    @staticmethod
+    def resolve_effective_true_cfg_scale(
+        req: OmniDiffusionRequest,
+        default_true_cfg_scale: float = 4.0,
+    ) -> float:
+        sampling_params = req.sampling_params
+        true_cfg_scale = sampling_params.true_cfg_scale
+        guidance_scale = sampling_params.guidance_scale
+        guidance_provided = sampling_params.guidance_scale_provided
+        if true_cfg_scale is None:
+            return float(guidance_scale if guidance_provided else default_true_cfg_scale)
+        if guidance_provided and not math.isclose(float(guidance_scale), float(true_cfg_scale)):
+            raise ValueError(
+                "JoyAI-Image-Edit treats `guidance_scale` as a compatibility alias for "
+                "`true_cfg_scale`. Provide only one value, or provide matching values."
+            )
+        return float(true_cfg_scale)

vllm_omni/diffusion/registry.py

@@ -51,6 +51,11 @@ _DIFFUSION_MODELS = {
         "pipeline_qwen_image_edit_plus",
         "QwenImageEditPlusPipeline",
     ),
+    "JoyImageEditPipeline": (
+        "joy_image",
+        "pipeline_joy_image_edit",
+        "JoyImageEditPipeline",
+    ),
     "QwenImageLayeredPipeline": (
         "qwen_image",
         "pipeline_qwen_image_layered",
@@ -660,6 +665,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = {
     "AnimaPipeline": "get_anima_post_process_func",
     "QwenImageEditPipeline": "get_qwen_image_edit_post_process_func",
     "QwenImageEditPlusPipeline": "get_qwen_image_edit_plus_post_process_func",
+    "JoyImageEditPipeline": "get_joy_image_edit_post_process_func",
     "QwenImage21Pipeline": "get_qwen_image_21_post_process_func",
     "GlmImagePipeline": "get_glm_image_post_process_func",
     "ZImagePipeline": "get_post_process_func",
@@ -750,6 +756,7 @@ _DIFFUSION_PRE_PROCESS_FUNCS = {
     "BooguImageTurboPipeline": "get_boogu_image_pre_process_func",
     "QwenImageEditPipeline": "get_qwen_image_edit_pre_process_func",
     "QwenImageEditPlusPipeline": "get_qwen_image_edit_plus_pre_process_func",
+    "JoyImageEditPipeline": "get_joy_image_edit_pre_process_func",
     "QwenImage21Pipeline": "get_qwen_image_21_pre_process_func",
     "LongCatImageEditPipeline": "get_longcat_image_edit_pre_process_func",
     "LongCatVideoAvatarPipeline": "get_longcat_video_avatar_pre_process_func",