diff --git a/.buildkite/common/ci_source_file_dependencies.yml b/.buildkite/common/ci_source_file_dependencies.yml index d7068cfca..c24ecd079 100644 --- a/.buildkite/common/ci_source_file_dependencies.yml +++ b/.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 diff --git a/.buildkite/cuda/test-nightly.yml b/.buildkite/cuda/test-nightly.yml index 86e794c7a..61b7c28a7 100644 --- a/.buildkite/cuda/test-nightly.yml +++ b/.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 diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index e16e0a5db..c5a2bf134 100644 --- a/docs/models/supported_models.md +++ b/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` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — | diff --git a/docs/user_guide/diffusion/cpu_offload.md b/docs/user_guide/diffusion/cpu_offload.md index bd34bf3bf..9e2d19d0e 100644 --- a/docs/user_guide/diffusion/cpu_offload.md +++ b/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` | diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md index 830b20597..0d76c7447 100644 --- a/docs/user_guide/diffusion_features.md +++ b/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)** | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | diff --git a/docs/user_guide/examples/offline_inference/image_to_image.md b/docs/user_guide/examples/offline_inference/image_to_image.md index b066eaee3..88adccbc1 100644 --- a/docs/user_guide/examples/offline_inference/image_to_image.md +++ b/docs/user_guide/examples/offline_inference/image_to_image.md @@ -3,7 +3,7 @@ Source . -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. diff --git a/docs/user_guide/examples/online_serving/image_to_image.md b/docs/user_guide/examples/online_serving/image_to_image.md index 89048ae71..cc2357547 100644 --- a/docs/user_guide/examples/online_serving/image_to_image.md +++ b/docs/user_guide/examples/online_serving/image_to_image.md @@ -3,10 +3,14 @@ Source . -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 diff --git a/examples/offline_inference/image_to_image/image_edit.py b/examples/offline_inference/image_to_image/image_edit.py index 4985117f6..bcced8bf4 100644 --- a/examples/offline_inference/image_to_image/image_edit.py +++ b/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 diff --git a/examples/offline_inference/image_to_image/image_to_image.md b/examples/offline_inference/image_to_image/image_to_image.md index 1c1a5ff3a..f738cedbd 100644 --- a/examples/offline_inference/image_to_image/image_to_image.md +++ b/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. diff --git a/examples/online_serving/image_to_image/README.md b/examples/online_serving/image_to_image/README.md index 8a6c6a0e6..28482c4a1 100644 --- a/examples/online_serving/image_to_image/README.md +++ b/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 diff --git a/recipes/JD/JoyAI-Image-Edit.md b/recipes/JD/JoyAI-Image-Edit.md new file mode 100644 index 000000000..23bbb55ac --- /dev/null +++ b/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. diff --git a/recipes/README.md b/recipes/README.md index 389d22024..b94cbf488 100644 --- a/recipes/README.md +++ b/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 | diff --git a/tests/diffusion/models/joy_image/test_joy_image.py b/tests/diffusion/models/joy_image/test_joy_image.py new file mode 100644 index 000000000..0bc582975 --- /dev/null +++ b/tests/diffusion/models/joy_image/test_joy_image.py @@ -0,0 +1,1636 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import json +import logging +from types import SimpleNamespace + +import pytest +import torch +import torch.nn.functional as F +from PIL import Image + +from vllm_omni.diffusion.attention.layer import Attention +from vllm_omni.diffusion.config import set_current_diffusion_config +from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec +from vllm_omni.diffusion.model_metadata import get_diffusion_model_metadata +from vllm_omni.diffusion.models.joy_image import ( + pipeline_joy_image_edit as joy_pipeline_module, +) +from vllm_omni.diffusion.models.joy_image.cfg_parallel import JoyImageEditCFGParallelMixin +from vllm_omni.diffusion.models.joy_image.joy_image_edit_transformer import ( + JoyImageAttention, + JoyImageEditTransformer3DModel, + _raise_if_unsupported_tensor_parallel, +) +from vllm_omni.diffusion.models.joy_image.pipeline_joy_image_edit import ( + JOY_MAX_IMAGE_SEQ_LEN, + JOY_PROMPT_TEMPLATE, + JOY_VISION_TOKEN, + JoyImageEditPipeline, + _cast_floating_model_inputs, + _format_qwen_multimodal_prompt, + _get_transformer_config_kwargs_from_od_config, + _raise_if_unsupported_hsdp, + _resize_center_crop, + _should_defer_component_device_placement, + _uses_model_level_cpu_offload, + get_joy_image_edit_pre_process_func, + retrieve_timesteps, +) +from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy +from vllm_omni.diffusion.offloader.layerwise_backend import LayerWiseOffloadBackend +from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery +from vllm_omni.diffusion.request import DUMMY_DIFFUSION_REQUEST_ID, OmniDiffusionRequest + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +def _torch_sdpa_diffusion_config(): + return SimpleNamespace( + diffusion_attention_config=AttentionConfig(default=AttentionSpec(backend="TORCH_SDPA")), + parallel_config=SimpleNamespace(ring_degree=1), + diffusion_kv_cache_dtype=None, + diffusion_kv_cache_skip_step_indices=None, + diffusion_kv_cache_skip_layer_indices=None, + ) + + +def _make_joy_attention(*, dtype: torch.dtype = torch.float32) -> JoyImageAttention: + with set_current_diffusion_config(_torch_sdpa_diffusion_config()): + attention = JoyImageAttention( + dim=32, + num_attention_heads=4, + attention_head_dim=8, + prefix="double_blocks.0.attn", + ) + return attention.to(dtype=dtype) + + +def _reference_joy_attention_output( + attention: JoyImageAttention, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + img_query, img_key, img_value = attention._stream_qkv( + attention.img_attn_qkv, + attention.img_attn_q_norm, + attention.img_attn_k_norm, + hidden_states, + ) + txt_query, txt_key, txt_value = attention._stream_qkv( + attention.txt_attn_qkv, + attention.txt_attn_q_norm, + attention.txt_attn_k_norm, + encoder_hidden_states, + ) + 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) + + if attention_mask is not None and attention_mask.ndim == 2: + attention_mask = attention_mask.to(torch.bool)[:, None, None, :] + + joint_hidden_states = F.scaled_dot_product_attention( + joint_query.transpose(1, 2), + joint_key.transpose(1, 2), + joint_value.transpose(1, 2), + attn_mask=attention_mask, + dropout_p=0.0, + is_causal=False, + scale=1.0 / (attention.head_dim**0.5), + ) + joint_hidden_states = joint_hidden_states.transpose(1, 2).flatten(2, 3).to(joint_query.dtype) + + image_seq_len = hidden_states.shape[1] + return ( + attention.img_attn_proj(joint_hidden_states[:, :image_seq_len]), + attention.txt_attn_proj(joint_hidden_states[:, image_seq_len:]), + ) + + +def _assert_attention_outputs_close( + actual: tuple[torch.Tensor, torch.Tensor], + expected: tuple[torch.Tensor, torch.Tensor], + *, + dtype: torch.dtype, +) -> None: + atol = rtol = 5e-3 if dtype is torch.bfloat16 else 1e-5 + max_abs = max( + (actual_item - expected_item).abs().max().item() for actual_item, expected_item in zip(actual, expected) + ) + assert max_abs <= atol + for actual_item, expected_item in zip(actual, expected): + torch.testing.assert_close(actual_item, expected_item, atol=atol, rtol=rtol) + + +def _write_model_configs(tmp_path): + (tmp_path / "vae").mkdir() + (tmp_path / "transformer").mkdir() + (tmp_path / "vae" / "config.json").write_text( + json.dumps( + { + "scale_factor_spatial": 8, + "z_dim": 16, + "latents_mean": [0.0] * 16, + "latents_std": [1.0] * 16, + } + ) + ) + (tmp_path / "transformer" / "config.json").write_text( + json.dumps( + { + "patch_size": [1, 2, 2], + "hidden_size": 32, + "num_attention_heads": 4, + "num_layers": 1, + "in_channels": 4, + "out_channels": 4, + "text_dim": 16, + } + ) + ) + return SimpleNamespace( + model=str(tmp_path), + model_class_name="JoyImageEditPipeline", + ) + + +def _make_params(**overrides): + guidance_scale = overrides.pop("guidance_scale", None) + values = { + "height": None, + "width": None, + "guidance_scale": guidance_scale if guidance_scale is not None else 1.0, + "guidance_scale_provided": guidance_scale is not None, + "true_cfg_scale": None, + "num_outputs_per_prompt": 1, + "num_inference_steps": None, + "timesteps": None, + "sigmas": None, + "max_sequence_length": None, + "latents": None, + "output_type": None, + "generator": None, + "seed": None, + "cfg_normalize": False, + "guidance_scale_2": None, + "do_classifier_free_guidance": False, + } + values.update(overrides) + return SimpleNamespace(**values) + + +def _make_request(*, prompt=None, params=None): + prompt = prompt or { + "prompt": "make it brighter", + "multi_modal_data": {"image": Image.new("RGB", (2048, 1024), color="white")}, + } + params = params or _make_params() + return SimpleNamespace( + prompts=[prompt], + sampling_params=params, + request_id="joy-test", + is_dummy_run=lambda: False, + ) + + +def _make_single_omni_request(*, prompt=None, params=None): + prompt = prompt or { + "prompt": "make it brighter", + "multi_modal_data": {"image": Image.new("RGB", (2048, 1024), color="white")}, + } + params = params or _make_params() + return OmniDiffusionRequest( + prompt=prompt, + sampling_params=params, + request_id="joy-test", + ) + + +class _TinyJoyLayerwisePipeline(torch.nn.Module): + _dit_modules = ["transformer"] + _encoder_modules: list[str] = [] + _vae_modules: list[str] = [] + _resident_modules: list[str] = [] + _offload_transformer_if_model_level = JoyImageEditPipeline._offload_transformer_if_model_level + + def __init__(self, *, device: torch.device, cleanup_mode: str): + super().__init__() + self.device = device + self.scheduler = SimpleNamespace( + set_timesteps=lambda num_inference_steps, device: setattr( + self.scheduler, + "timesteps", + torch.tensor([1.0], device=device), + ) + ) + self.latent_channels = 4 + self.od_config = SimpleNamespace( + enable_cpu_offload=False, + enable_layerwise_offload=True, + pin_cpu_memory=False, + parallel_config=SimpleNamespace(use_hsdp=False), + ) + self.transformer = JoyImageEditTransformer3DModel( + in_channels=4, + out_channels=4, + hidden_size=32, + text_dim=16, + num_layers=2, + num_attention_heads=4, + patch_size=(1, 2, 2), + ).to(device=device, dtype=torch.bfloat16) + self.vae = torch.nn.Linear(1, 1) + self.cleanup_mode = cleanup_mode + + def encode_prompt(self, *args, **kwargs): + return ( + torch.zeros(1, 5, 16, device=self.device, dtype=torch.bfloat16), + torch.ones(1, 5, dtype=torch.long, device=self.device), + ) + + def resolve_effective_true_cfg_scale(self, req, default_true_cfg_scale=4.0): + return JoyImageEditPipeline.resolve_effective_true_cfg_scale( + req, + default_true_cfg_scale=default_true_cfg_scale, + ) + + def _prepare_latents(self, **kwargs): + latents = torch.randn(1, 2, 4, 1, 4, 4, device=self.device, dtype=torch.bfloat16) + return latents, latents[:, :1].clone() + + def diffuse(self, **kwargs): + return self.transformer( + hidden_states=kwargs["latents"], + timestep=kwargs["timesteps"][0], + encoder_hidden_states=kwargs["prompt_embeds"], + encoder_hidden_states_mask=kwargs["prompt_embeds_mask"], + return_dict=False, + )[0] + + def _decode_latents(self, latents): + JoyImageEditPipeline._prepare_vae_for_decode(self) + return latents[:, -1, :, 0] + + +def _make_layerwise_forward_request(*, dummy: bool = False): + request = _make_request( + prompt={ + "prompt": "make it brighter", + "additional_information": { + "image_tensor": torch.zeros(1, 3, 1, 16, 16), + "prompt_image": Image.new("RGB", (16, 16)), + "height": 16, + "width": 16, + }, + }, + params=_make_params(guidance_scale=1.0, num_inference_steps=1), + ) + request.is_dummy_run = lambda: dummy + return request + + +def _run_tiny_joy_layerwise_sequence(first_mode: str): + device = torch.device("cuda:0") + pipeline = _TinyJoyLayerwisePipeline(device=device, cleanup_mode=first_mode) + backend = LayerWiseOffloadBackend( + OffloadConfig(strategy=OffloadStrategy.LAYER_WISE, pin_cpu_memory=False), + device, + ) + backend.enable(pipeline) + try: + if first_mode == "dummy": + first_request = _make_layerwise_forward_request(dummy=True) + first_output_type = "pil" + elif first_mode == "latent": + first_request = _make_layerwise_forward_request(dummy=False) + first_output_type = "latent" + elif first_mode == "decoded": + first_request = _make_layerwise_forward_request(dummy=False) + first_output_type = "pil" + else: + raise AssertionError(f"unexpected first mode {first_mode!r}") + + with torch.inference_mode(): + JoyImageEditPipeline.forward(pipeline, first_request, output_type=first_output_type) + + # These non-block modules execute before double_blocks[0], so layerwise + # hooks cannot repair them lazily after request cleanup. + assert next(pipeline.transformer.img_in.parameters()).device == device + assert next(pipeline.transformer.condition_embedder.parameters()).device == device + assert next(pipeline.transformer.proj_out.parameters()).device == device + + second_request = _make_layerwise_forward_request(dummy=False) + with torch.inference_mode(): + output = JoyImageEditPipeline.forward(pipeline, second_request, output_type="pil") + finally: + backend.disable() + torch.accelerator.empty_cache() + + assert output.output.shape == (1, 4, 4, 4) + + +def test_joy_registry_and_metadata_entries(tmp_path): + from vllm_omni.diffusion.registry import ( + DiffusionModelRegistry, + get_diffusion_post_process_func, + get_diffusion_pre_process_func, + ) + + od_config = _write_model_configs(tmp_path) + model_class = DiffusionModelRegistry._try_load_model_cls("JoyImageEditPipeline") + assert model_class is JoyImageEditPipeline + assert callable(get_diffusion_pre_process_func(od_config)) + postprocess = get_diffusion_post_process_func(od_config) + assert callable(postprocess) + assert len(postprocess(torch.zeros(1, 3, 2, 2))) == 1 + + metadata = get_diffusion_model_metadata("JoyImageEditPipeline") + assert metadata.supports_multimodal_inputs is True + assert metadata.max_multimodal_image_inputs == 1 + + +def test_get_model_path_downloads_only_runtime_configs(monkeypatch): + captured = {} + + def fake_download(model_name, revision, allow_patterns, **kwargs): + captured.update( + { + "model_name": model_name, + "revision": revision, + "allow_patterns": allow_patterns, + "kwargs": kwargs, + } + ) + return "/tmp/joy-configs" + + monkeypatch.setattr( + joy_pipeline_module, + "download_weights_from_hf_specific", + fake_download, + ) + + assert joy_pipeline_module._get_model_path("repo/joy") == "/tmp/joy-configs" + assert captured == { + "model_name": "repo/joy", + "revision": None, + "allow_patterns": ["vae/config.json", "transformer/config.json"], + "kwargs": {"require_all": True}, + } + + +def test_transformer_config_kwargs_from_od_config_filters_metadata(): + class FakeTransformerConfig: + def to_dict(self): + return { + "_class_name": "JoyImageEditTransformer3DModel", + "_diffusers_version": "0.38.0", + "hidden_size": 32, + "num_layers": 1, + "num_attention_heads": 4, + } + + od_config = SimpleNamespace(tf_model_config=FakeTransformerConfig()) + + assert _get_transformer_config_kwargs_from_od_config(od_config) == { + "hidden_size": 32, + "num_layers": 1, + "num_attention_heads": 4, + } + + +def test_defer_component_device_placement_for_offload_only(): + assert ( + _should_defer_component_device_placement( + SimpleNamespace( + enable_cpu_offload=False, + enable_layerwise_offload=False, + parallel_config=SimpleNamespace(use_hsdp=False), + ) + ) + is False + ) + assert _should_defer_component_device_placement(SimpleNamespace(enable_cpu_offload=True)) is True + assert _should_defer_component_device_placement(SimpleNamespace(enable_layerwise_offload=True)) is True + assert ( + _should_defer_component_device_placement(SimpleNamespace(parallel_config=SimpleNamespace(use_hsdp=True))) + is False + ) + + +@pytest.mark.parametrize( + ("enable_cpu_offload", "enable_layerwise_offload", "expected"), + [ + (True, False, True), + (False, True, False), + (True, True, False), + (False, False, False), + ], +) +def test_model_level_cpu_offload_predicate_matches_effective_strategy( + enable_cpu_offload, + enable_layerwise_offload, + expected, +): + od_config = SimpleNamespace( + enable_cpu_offload=enable_cpu_offload, + enable_layerwise_offload=enable_layerwise_offload, + ) + + assert _uses_model_level_cpu_offload(od_config) is expected + + +def test_model_level_cpu_offload_cleanup_still_moves_transformer(monkeypatch): + pipeline = object.__new__(JoyImageEditPipeline) + torch.nn.Module.__init__(pipeline) + pipeline.od_config = SimpleNamespace( + enable_cpu_offload=True, + enable_layerwise_offload=False, + pin_cpu_memory=False, + ) + pipeline.transformer = torch.nn.Linear(1, 1) + move_calls = [] + + def fake_move_params(module, target_device, **kwargs): + move_calls.append((module, target_device, kwargs)) + + monkeypatch.setattr( + joy_pipeline_module.SequentialOffloadHook, + "_move_params", + fake_move_params, + ) + monkeypatch.setattr(joy_pipeline_module.current_omni_platform, "empty_cache", lambda: None) + + JoyImageEditPipeline._offload_transformer_if_model_level(pipeline) + + assert move_calls == [ + ( + pipeline.transformer, + torch.device("cpu"), + {"non_blocking": False, "pin_memory": False}, + ) + ] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("first_mode", ["dummy", "decoded", "latent"]) +def test_layerwise_joy_lifecycle_survives_second_forward(first_mode): + _run_tiny_joy_layerwise_sequence(first_mode) + + +def test_joy_hsdp_is_explicitly_unsupported(): + with pytest.raises(ValueError, match="does not support HSDP"): + _raise_if_unsupported_hsdp(SimpleNamespace(parallel_config=SimpleNamespace(use_hsdp=True))) + + _raise_if_unsupported_hsdp(SimpleNamespace(parallel_config=SimpleNamespace(use_hsdp=False))) + _raise_if_unsupported_hsdp(SimpleNamespace()) + + +def test_joy_tensor_parallelism_is_explicitly_unsupported(): + with pytest.raises(ValueError, match="do not support Tensor Parallelism"): + _raise_if_unsupported_tensor_parallel(SimpleNamespace(parallel_config=SimpleNamespace(tensor_parallel_size=2))) + + _raise_if_unsupported_tensor_parallel(SimpleNamespace(parallel_config=SimpleNamespace(tensor_parallel_size=1))) + _raise_if_unsupported_tensor_parallel(SimpleNamespace(parallel_config=SimpleNamespace())) + _raise_if_unsupported_tensor_parallel(SimpleNamespace()) + _raise_if_unsupported_tensor_parallel(None) + + +def test_transformer_rejects_tensor_parallelism_before_building_layers(): + with pytest.raises(ValueError, match="do not support Tensor Parallelism"): + JoyImageEditTransformer3DModel( + od_config=SimpleNamespace(parallel_config=SimpleNamespace(tensor_parallel_size=2)), + in_channels=4, + out_channels=4, + hidden_size=32, + text_dim=16, + num_layers=1, + num_attention_heads=4, + patch_size=(1, 2, 2), + ) + + +def test_pipeline_rejects_tensor_parallelism_before_loading_components(): + with pytest.raises(ValueError, match="do not support Tensor Parallelism"): + JoyImageEditPipeline( + od_config=SimpleNamespace(parallel_config=SimpleNamespace(use_hsdp=False, tensor_parallel_size=2)) + ) + + +def test_component_discovery_treats_vae_as_offload_peer(): + pipeline = object.__new__(JoyImageEditPipeline) + torch.nn.Module.__init__(pipeline) + pipeline.transformer = torch.nn.Linear(1, 1) + pipeline.text_encoder = torch.nn.Linear(1, 1) + pipeline.vae = torch.nn.Linear(1, 1) + + modules = ModuleDiscovery.discover(pipeline) + + assert modules.dit_names == ["transformer"] + assert modules.encoder_names == ["text_encoder", "vae"] + assert modules.vaes == [] + + +def test_preprocess_requires_exactly_one_image(tmp_path): + od_config = _write_model_configs(tmp_path) + preprocess = get_joy_image_edit_pre_process_func(od_config) + no_image = _make_request(prompt={"prompt": "x", "multi_modal_data": {}}) + with pytest.raises(ValueError, match="exactly one input image"): + preprocess(no_image) + + multi_image = _make_request( + prompt={ + "prompt": "x", + "multi_modal_data": {"image": [Image.new("RGB", (32, 32)), Image.new("RGB", (32, 32))]}, + } + ) + with pytest.raises(ValueError, match="exactly one image"): + preprocess(multi_image) + + +def test_preprocess_rejects_batched_prompts_and_partial_size(tmp_path): + od_config = _write_model_configs(tmp_path) + preprocess = get_joy_image_edit_pre_process_func(od_config) + batched = _make_request() + batched.prompts.append(batched.prompts[0]) + with pytest.raises(ValueError, match="exactly one prompt"): + preprocess(batched) + + partial_size = _make_request(params=_make_params(height=512, width=None)) + with pytest.raises(ValueError, match="both `height` and `width`"): + preprocess(partial_size) + + +def test_preprocess_uses_diffusers_bucket_for_default_image_size(tmp_path): + od_config = _write_model_configs(tmp_path) + preprocess = get_joy_image_edit_pre_process_func(od_config) + + request = preprocess(_make_request()) + prompt = request.prompts[0] + info = prompt["additional_information"] + height = info["height"] + width = info["width"] + + assert (height, width) == (704, 1408) + assert (height // 16) * (width // 16) <= JOY_MAX_IMAGE_SEQ_LEN + assert info["image_tensor"].shape == (1, 3, 1, height, width) + assert info["original_size"] == (2048, 1024) + assert info["resized_size"] == (width, height) + assert request.sampling_params.height == height + assert request.sampling_params.width == width + + +def test_preprocess_accepts_single_omni_request_prompt(tmp_path): + od_config = _write_model_configs(tmp_path) + preprocess = get_joy_image_edit_pre_process_func(od_config) + + request = preprocess(_make_single_omni_request()) + prompt = request.prompt + info = prompt["additional_information"] + + assert isinstance(prompt, dict) + assert not hasattr(request, "prompts") + assert info["image_tensor"].shape == (1, 3, 1, 704, 1408) + assert request.sampling_params.height == 704 + assert request.sampling_params.width == 1408 + + +def test_resize_center_crop_crops_instead_of_stretching(): + image = Image.new("RGB", (4, 2)) + pixels = image.load() + colors = [ + (255, 0, 0), + (0, 255, 0), + (0, 0, 255), + (255, 255, 255), + ] + for x, color in enumerate(colors): + for y in range(2): + pixels[x, y] = color + + cropped = _resize_center_crop(image, height=2, width=2) + cropped_pixels = cropped.load() + + assert cropped.size == (2, 2) + assert [[cropped_pixels[x, y] for x in range(2)] for y in range(2)] == [ + [(0, 255, 0), (0, 0, 255)], + [(0, 255, 0), (0, 0, 255)], + ] + + +def test_joy_cfg_normalize_uses_channel_dimension(): + mixin = JoyImageEditCFGParallelMixin() + noise_pred = torch.tensor([[[[[[3.0, 4.0]]], [[[4.0, 3.0]]]]]]) + comb_pred = torch.tensor([[[[[[6.0, 8.0]]], [[[8.0, 6.0]]]]]]) + + normalized = mixin.cfg_normalize_function(noise_pred, comb_pred) + expected = comb_pred * ( + torch.norm(noise_pred, dim=2, keepdim=True) / torch.norm(comb_pred, dim=2, keepdim=True).clamp_min(1e-6) + ) + + assert torch.equal(normalized, expected) + + +def test_joy_diffuse_preserves_scheduler_timestep_dtype_for_time_embedding(): + class FakeProgressBar: + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_value, traceback): + return False + + def update(self): + pass + + class FakeScheduler: + def set_begin_index(self, index): + self.begin_index = index + + class FakeMixin(JoyImageEditCFGParallelMixin): + def __init__(self): + self.scheduler = FakeScheduler() + self.captured_timestep_dtype = None + self._interrupt = False + + @property + def interrupt(self): + return self._interrupt + + def progress_bar(self, total): + return FakeProgressBar() + + def predict_noise_maybe_with_cfg(self, **kwargs): + self.captured_timestep_dtype = kwargs["positive_kwargs"]["timestep"].dtype + return torch.zeros_like(kwargs["positive_kwargs"]["hidden_states"]) + + def scheduler_step_maybe_with_cfg(self, noise_pred, timestep, latents, do_true_cfg): + return latents + + mixin = FakeMixin() + latents = torch.zeros(1, 2, 4, 1, 4, 4, dtype=torch.bfloat16) + image_latents = torch.zeros(1, 1, 4, 1, 4, 4, dtype=torch.bfloat16) + + mixin.diffuse( + latents=latents, + image_latents=image_latents, + prompt_embeds=torch.zeros(1, 2, 4, dtype=torch.bfloat16), + prompt_embeds_mask=torch.ones(1, 2, dtype=torch.long), + negative_prompt_embeds=None, + negative_prompt_embeds_mask=None, + timesteps=torch.tensor([999.125], dtype=torch.float32), + do_true_cfg=False, + true_cfg_scale=1.0, + ) + + assert mixin.captured_timestep_dtype == torch.float32 + + +def test_joy_diffuse_replaces_reference_latents_out_of_place(): + class FakeProgressBar: + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_value, traceback): + return False + + def update(self): + pass + + class FakeScheduler: + def set_begin_index(self, index): + self.begin_index = index + + class FakeMixin(JoyImageEditCFGParallelMixin): + def __init__(self): + self.scheduler = FakeScheduler() + self.scheduler_returned_latents = None + self._interrupt = False + + @property + def interrupt(self): + return self._interrupt + + def progress_bar(self, total): + return FakeProgressBar() + + def predict_noise_maybe_with_cfg(self, **kwargs): + return torch.zeros_like(kwargs["positive_kwargs"]["hidden_states"]) + + def scheduler_step_maybe_with_cfg(self, noise_pred, timestep, latents, do_true_cfg): + stepped = latents.clone() + stepped[:, :1] = -5 + self.scheduler_returned_latents = stepped + return stepped + + mixin = FakeMixin() + latents = torch.arange(3, dtype=torch.float32).view(1, 3, 1, 1, 1, 1) + original_latents = latents.clone() + image_latents = torch.full((1, 1, 1, 1, 1, 1), 100.0) + + output = mixin.diffuse( + latents=latents, + image_latents=image_latents, + prompt_embeds=torch.zeros(1, 2, 4), + prompt_embeds_mask=torch.ones(1, 2, dtype=torch.long), + negative_prompt_embeds=None, + negative_prompt_embeds_mask=None, + timesteps=torch.tensor([999.125], dtype=torch.float32), + do_true_cfg=False, + true_cfg_scale=1.0, + ) + + assert torch.equal(latents, original_latents) + assert mixin.scheduler_returned_latents is not None + assert torch.equal(mixin.scheduler_returned_latents[:, :1], torch.full_like(image_latents, -5)) + assert torch.equal(output[:, :1], image_latents) + + +def test_format_qwen_multimodal_prompt_matches_diffusers_image_placeholder_replacement(): + formatted = _format_qwen_multimodal_prompt("make it brighter") + + assert formatted == (f"<|im_start|>user\n{JOY_VISION_TOKEN}make it brighter<|im_end|>\n") + + +def test_preprocess_maps_explicit_size_to_nearest_diffusers_bucket(tmp_path): + od_config = _write_model_configs(tmp_path) + preprocess = get_joy_image_edit_pre_process_func(od_config) + params = _make_params(height=2048, width=2048) + + request = preprocess(_make_request(params=params)) + + assert request.sampling_params.height == 1024 + assert request.sampling_params.width == 1024 + + +@pytest.mark.parametrize( + "params_kwargs, expected_scale", + [ + ({}, 4.0), + ({"guidance_scale": 3.5}, 3.5), + ({"true_cfg_scale": 4.0}, 4.0), + ({"guidance_scale": 4.0, "true_cfg_scale": 4.0}, 4.0), + ], +) +def test_resolve_effective_true_cfg_scale(params_kwargs, expected_scale): + request = _make_request(params=_make_params(**params_kwargs)) + assert JoyImageEditPipeline.resolve_effective_true_cfg_scale(request) == expected_scale + + +@pytest.mark.parametrize("guidance_scale", [1.0, 3.0]) +def test_explicit_guidance_scale_conflicts_with_true_cfg_scale(guidance_scale): + request = _make_request(params=_make_params(guidance_scale=guidance_scale, true_cfg_scale=4.0)) + with pytest.raises(ValueError, match="compatibility alias"): + JoyImageEditPipeline.resolve_effective_true_cfg_scale(request) + + +def test_pad_prompt_embeds_keeps_last_tokens_and_builds_mask(): + first = torch.arange(20, dtype=torch.float32).reshape(5, 4) + second = torch.arange(8, dtype=torch.float32).reshape(2, 4) + + prompt_embeds, prompt_mask = JoyImageEditPipeline._pad_prompt_embeds( + [first, second], + max_sequence_length=3, + ) + + assert prompt_embeds.shape == (2, 3, 4) + assert torch.equal(prompt_embeds[0], first[-3:]) + assert torch.equal(prompt_embeds[1, :2], second) + assert torch.equal(prompt_embeds[1, 2], torch.zeros(4)) + assert torch.equal( + prompt_mask, + torch.tensor([[1, 1, 1], [1, 1, 0]], dtype=torch.long), + ) + + with pytest.raises(ValueError, match="greater than 0"): + JoyImageEditPipeline._pad_prompt_embeds([first], max_sequence_length=0) + + +@pytest.mark.parametrize( + ("prompt_embeds", "prompt_embeds_mask", "match"), + [ + ( + torch.zeros(2, 4), + torch.ones(2, 4, dtype=torch.long), + "prompt_embeds must be a 3D tensor", + ), + ( + torch.zeros(2, 4, 8), + torch.ones(2, 4, 1, dtype=torch.long), + "prompt_embeds_mask must be a 2D tensor", + ), + ( + torch.zeros(2, 4, 8), + torch.ones(3, 4, dtype=torch.long), + "same batch size", + ), + ( + torch.zeros(2, 4, 8), + torch.ones(2, 5, dtype=torch.long), + "same sequence length", + ), + ], +) +def test_encode_prompt_validates_precomputed_embeds_and_mask_shapes( + prompt_embeds, + prompt_embeds_mask, + match, +): + pipeline = object.__new__(JoyImageEditPipeline) + + with pytest.raises(ValueError, match=match): + JoyImageEditPipeline.encode_prompt( + pipeline, + "prompt", + Image.new("RGB", (16, 16)), + num_images_per_prompt=1, + prompt_embeds=prompt_embeds, + prompt_embeds_mask=prompt_embeds_mask, + ) + + +def test_encode_prompt_validation_uses_negative_parameter_names(): + pipeline = object.__new__(JoyImageEditPipeline) + + with pytest.raises(ValueError, match="negative_prompt_embeds_mask must be a 2D tensor"): + JoyImageEditPipeline.encode_prompt( + pipeline, + "", + Image.new("RGB", (16, 16)), + num_images_per_prompt=1, + prompt_embeds=torch.zeros(1, 4, 8), + prompt_embeds_mask=torch.ones(1, 4, 1, dtype=torch.long), + embeds_name="negative_prompt_embeds", + mask_name="negative_prompt_embeds_mask", + ) + + +def test_last_layer_capture_passes_mm_token_type_ids(): + class FakeDecoderLayer(torch.nn.Module): + def forward(self, hidden_states): + return (hidden_states + 1.0,) + + class FakeLanguageModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.layers = torch.nn.ModuleList([FakeDecoderLayer()]) + + class FakeTextEncoder(torch.nn.Module): + def __init__(self): + super().__init__() + self.model = SimpleNamespace(language_model=FakeLanguageModel()) + self.calls = [] + + def forward(self, **kwargs): + self.calls.append(kwargs) + hidden_states = torch.zeros(1, 2, 4) + self.model.language_model.layers[-1](hidden_states) + + pipeline = object.__new__(JoyImageEditPipeline) + torch.nn.Module.__init__(pipeline) + pipeline.text_encoder = FakeTextEncoder() + mm_token_type_ids = torch.tensor([[0, 1]], dtype=torch.long) + model_inputs = SimpleNamespace( + input_ids=torch.tensor([[1, 2]], dtype=torch.long), + attention_mask=torch.tensor([[1, 1]], dtype=torch.long), + pixel_values=torch.zeros(1, 3, 16, 16), + image_grid_thw=torch.tensor([[1, 1, 1]], dtype=torch.long), + mm_token_type_ids=mm_token_type_ids, + extra_processor_field=torch.tensor([7], dtype=torch.long), + ) + + hidden_states = JoyImageEditPipeline._get_last_layer_pre_norm_hidden( + pipeline, + model_inputs, + ) + + assert torch.equal(hidden_states, torch.ones(1, 2, 4)) + assert pipeline.text_encoder.calls[0]["mm_token_type_ids"] is mm_token_type_ids + assert torch.equal(pipeline.text_encoder.calls[0]["extra_processor_field"], torch.tensor([7], dtype=torch.long)) + assert pipeline.text_encoder.calls[0]["output_hidden_states"] is False + + +def test_cast_floating_model_inputs_only_casts_pixel_tensors(): + class FakeModelInputs(dict): + def __getattr__(self, name): + return self[name] + + model_inputs = FakeModelInputs( + input_ids=torch.tensor([[1, 2]], dtype=torch.long), + attention_mask=torch.tensor([[1, 1]], dtype=torch.long), + pixel_values=torch.zeros(1, 3, dtype=torch.float32), + image_grid_thw=torch.tensor([[1, 1, 1]], dtype=torch.long), + ) + + _cast_floating_model_inputs(model_inputs, torch.bfloat16) + + assert model_inputs.input_ids.dtype == torch.long + assert model_inputs.attention_mask.dtype == torch.long + assert model_inputs.pixel_values.dtype == torch.bfloat16 + assert model_inputs.image_grid_thw.dtype == torch.long + + +def test_qwen_prompt_embeds_casts_processor_pixel_values_to_encoder_dtype(): + class FakeModelInputs(dict): + def __init__(self): + super().__init__( + input_ids=torch.tensor([[1, 2, 3]], dtype=torch.long), + attention_mask=torch.tensor([[1, 1, 1]], dtype=torch.long), + pixel_values=torch.zeros(1, 3, dtype=torch.float32), + image_grid_thw=torch.tensor([[1, 1, 1]], dtype=torch.long), + ) + + def __getattr__(self, name): + return self[name] + + def __setattr__(self, name, value): + self[name] = value + + def to(self, device): + for key, value in list(self.items()): + if isinstance(value, torch.Tensor): + self[key] = value.to(device) + return self + + class FakeProcessor: + def __init__(self): + self.model_inputs = FakeModelInputs() + + def __call__(self, **kwargs): + self.kwargs = kwargs + return self.model_inputs + + class FakeTextEncoder: + dtype = torch.bfloat16 + + pipeline = object.__new__(JoyImageEditPipeline) + torch.nn.Module.__init__(pipeline) + pipeline.device = torch.device("cpu") + pipeline.processor = FakeProcessor() + pipeline.text_encoder = FakeTextEncoder() + pipeline.prompt_template_encode = JOY_PROMPT_TEMPLATE + pipeline.prompt_template_encode_start_idx = 0 + pipeline.tokenizer_max_length = 8 + captured = {} + + def fake_get_hidden(model_inputs): + captured["pixel_dtype"] = model_inputs.pixel_values.dtype + captured["attention_mask_dtype"] = model_inputs.attention_mask.dtype + return torch.ones(1, 3, 4, dtype=torch.float32) + + pipeline._get_last_layer_pre_norm_hidden = fake_get_hidden + + prompt_embeds, prompt_mask = JoyImageEditPipeline._get_qwen_prompt_embeds( + pipeline, + "make it brighter", + Image.new("RGB", (16, 16)), + ) + + assert captured == { + "pixel_dtype": torch.bfloat16, + "attention_mask_dtype": torch.long, + } + text = pipeline.processor.kwargs["text"][0] + assert text.startswith("<|im_start|>system\n \\nDescribe") + assert "Describe the image by detailing the color, shape, size, texture" in text + assert (f"<|im_start|>user\n{JOY_VISION_TOKEN}make it brighter<|im_end|>\n") in text + assert text.endswith("<|im_start|>assistant\n") + assert prompt_embeds.dtype == torch.bfloat16 + assert prompt_mask.dtype == torch.long + + +def test_prepare_latents_casts_vae_encoded_latents_to_requested_dtype(): + pipeline = object.__new__(JoyImageEditPipeline) + torch.nn.Module.__init__(pipeline) + pipeline.latent_channels = 4 + pipeline.vae_scale_factor = 8 + + def fake_encode(image, generator): + assert image.dtype == torch.bfloat16 + return torch.ones(1, 4, 1, 8, 8, dtype=torch.float32) + + pipeline._encode_vae_image = fake_encode + + latents, image_latents = JoyImageEditPipeline._prepare_latents( + pipeline, + image=torch.zeros(1, 3, 1, 64, 64, dtype=torch.float32), + batch_size=1, + num_channels_latents=4, + height=64, + width=64, + dtype=torch.bfloat16, + device=torch.device("cpu"), + generator=None, + ) + + assert latents.dtype == torch.bfloat16 + assert image_latents.dtype == torch.bfloat16 + + +def test_prepare_latents_prepares_offloaded_vae_before_encode(): + pipeline = object.__new__(JoyImageEditPipeline) + pipeline.vae_scale_factor = 8 + calls = [] + pipeline._prepare_vae_for_decode = lambda: calls.append("prepare") + + def encode(image, generator): + calls.append("encode") + return torch.zeros(1, 4, 1, 2, 2, device=image.device, dtype=image.dtype) + + pipeline._encode_vae_image = encode + pipeline._prepare_latents( + image=torch.zeros(1, 3, 1, 16, 16), + batch_size=1, + num_channels_latents=4, + height=16, + width=16, + dtype=torch.float32, + device=torch.device("cpu"), + generator=None, + ) + + assert calls == ["prepare", "encode"] + + +def test_encode_vae_image_does_not_forward_noise_generator_to_posterior_sample(): + class FakeLatentDist: + def __init__(self): + self.received_generator = object() + + def sample(self, generator=None): + self.received_generator = generator + return torch.zeros(1, 4, 1, 8, 8) + + class FakeVAE: + def __init__(self): + self.config = SimpleNamespace( + latents_mean=[0.0, 0.0, 0.0, 0.0], + latents_std=[1.0, 1.0, 1.0, 1.0], + ) + self.latent_dist = FakeLatentDist() + + def encode(self, image): + return SimpleNamespace(latent_dist=self.latent_dist) + + pipeline = object.__new__(JoyImageEditPipeline) + torch.nn.Module.__init__(pipeline) + pipeline.vae = FakeVAE() + pipeline.latent_channels = 4 + + generator = torch.Generator(device="cpu").manual_seed(123) + JoyImageEditPipeline._encode_vae_image( + pipeline, + torch.zeros(1, 3, 1, 64, 64), + generator, + ) + + assert pipeline.vae.latent_dist.received_generator is None + + +@pytest.mark.parametrize( + "dtype", + [ + torch.float32, + torch.bfloat16, + ], +) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_joy_attention_native_sdpa_matches_torch_sdpa_without_mask(dtype): + torch.manual_seed(123) + device = torch.device("cuda") + attention = _make_joy_attention(dtype=dtype).to(device=device) + hidden_states = torch.randn(2, 3, 32, device=device, dtype=dtype) + encoder_hidden_states = torch.randn(2, 5, 32, device=device, dtype=dtype) + + actual = attention(hidden_states, encoder_hidden_states) + expected = _reference_joy_attention_output(attention, hidden_states, encoder_hidden_states) + + assert isinstance(attention.attn, Attention) + assert attention.attn.role == "joy_image.joint" + assert attention.attn.role_category == "self" + assert attention.attn.qkv_layout == "BSND" + _assert_attention_outputs_close(actual, expected, dtype=dtype) + + +@pytest.mark.parametrize( + "dtype", + [ + torch.float32, + torch.bfloat16, + ], +) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_joy_attention_native_sdpa_matches_torch_sdpa_with_padding_mask(dtype): + torch.manual_seed(456) + device = torch.device("cuda") + attention = _make_joy_attention(dtype=dtype).to(device=device) + hidden_states = torch.randn(2, 3, 32, device=device, dtype=dtype) + encoder_hidden_states = torch.randn(2, 5, 32, device=device, dtype=dtype) + attention_mask = torch.tensor( + [ + [1, 1, 1, 1, 1, 1, 1, 1], + [1, 1, 1, 1, 1, 1, 0, 0], + ], + device=device, + dtype=torch.bool, + ) + + actual = attention(hidden_states, encoder_hidden_states, attention_mask=attention_mask) + expected = _reference_joy_attention_output( + attention, + hidden_states, + encoder_hidden_states, + attention_mask=attention_mask, + ) + + _assert_attention_outputs_close(actual, expected, dtype=dtype) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_transformer_shape_and_masked_forward(): + device = torch.device("cuda") + transformer = JoyImageEditTransformer3DModel( + in_channels=4, + out_channels=4, + hidden_size=32, + text_dim=16, + num_layers=1, + num_attention_heads=4, + patch_size=(1, 2, 2), + ).to(device=device, dtype=torch.bfloat16) + hidden_states = torch.randn(2, 2, 4, 1, 4, 4, device=device, dtype=torch.bfloat16) + encoder_hidden_states = torch.randn(2, 5, 16, device=device, dtype=torch.bfloat16) + encoder_hidden_states_mask = torch.tensor([[1, 1, 1, 1, 1], [1, 1, 1, 0, 0]], device=device) + + output = transformer( + hidden_states=hidden_states, + timestep=torch.tensor([1.0, 2.0], device=device), + encoder_hidden_states=encoder_hidden_states, + encoder_hidden_states_mask=encoder_hidden_states_mask, + return_dict=False, + )[0] + + assert output.shape == hidden_states.shape + assert isinstance(transformer.double_blocks[0].attn.attn, Attention) + + +def test_transformer_from_config_file_loads_checkpoint_config(tmp_path): + _write_model_configs(tmp_path) + + transformer = JoyImageEditTransformer3DModel.from_config_file(tmp_path / "transformer" / "config.json") + + assert transformer.patch_size == (1, 2, 2) + assert transformer.hidden_size == 32 + assert transformer.num_attention_heads == 4 + assert transformer.num_layers == 1 + assert transformer.text_dim == 16 + assert transformer.theta == 256 + assert len(transformer.double_blocks) == 1 + + +def test_transformer_load_weights_maps_diffusers_prefix(): + transformer = JoyImageEditTransformer3DModel( + in_channels=4, + out_channels=4, + hidden_size=32, + text_dim=16, + num_layers=1, + num_attention_heads=4, + patch_size=(1, 2, 2), + ) + weight = torch.randn_like(transformer.img_in.weight) + + loaded = transformer.load_weights([("transformer.patch_embedding.weight", weight)]) + + assert loaded == {"img_in.weight"} + assert torch.equal(transformer.img_in.weight, weight) + + text_weight = torch.randn_like(transformer.condition_embedder.text_embedder.linear_1.weight) + loaded = transformer.load_weights([("transformer.condition_embedder.text_embedder.linear_1.weight", text_weight)]) + + assert loaded == {"condition_embedder.text_embedder.linear_1.weight"} + assert torch.equal(transformer.condition_embedder.text_embedder.linear_1.weight, text_weight) + + +def test_transformer_load_weights_warns_for_unknown_weight(caplog): + transformer = JoyImageEditTransformer3DModel( + in_channels=4, + out_channels=4, + hidden_size=32, + text_dim=16, + num_layers=1, + num_attention_heads=4, + patch_size=(1, 2, 2), + ) + + caplog.set_level(logging.WARNING) + loaded = transformer.load_weights([("transformer.unmapped.weight", torch.empty(1))]) + + assert loaded == set() + assert "Skipping JoyAI-Image-Edit transformer weight transformer.unmapped.weight" in caplog.text + assert "checkpoint mismatch" in caplog.text + + +def test_latent_normalization_round_trip(): + pipeline = object.__new__(JoyImageEditPipeline) + pipeline.latent_channels = 4 + pipeline.vae = SimpleNamespace( + config=SimpleNamespace( + latents_mean=[0.1, -0.2, 0.3, -0.4], + latents_std=[0.5, 0.75, 1.25, 2.0], + ) + ) + latents = torch.randn(2, 4, 1, 2, 2) + + latents_mean, latents_std = JoyImageEditPipeline._latent_stats( + pipeline, + latents.device, + latents.dtype, + ) + normalized = (latents - latents_mean) / latents_std + + assert torch.allclose(normalized * latents_std + latents_mean, latents) + + +def test_decode_latents_unnormalizes_and_selects_target_slot(): + class FakeVAE: + config = SimpleNamespace( + latents_mean=[0.1, -0.2, 0.3, -0.4], + latents_std=[0.5, 0.75, 1.25, 2.0], + ) + + def decode(self, latents, return_dict=False): + return (latents,) + + pipeline = object.__new__(JoyImageEditPipeline) + pipeline.latent_channels = 4 + pipeline.vae = FakeVAE() + latents = torch.zeros(1, 2, 4, 1, 2, 2) + latents[:, -1] = 2.0 + + decoded = JoyImageEditPipeline._decode_latents(pipeline, latents) + + latents_mean, latents_std = JoyImageEditPipeline._latent_stats( + pipeline, + latents.device, + latents.dtype, + ) + assert decoded.shape == (1, 4, 2, 2) + assert torch.equal(decoded, (latents[:, -1] * latents_std + latents_mean)[:, :, 0]) + + +def test_prepare_latents_stacks_reference_first_target_last(): + pipeline = object.__new__(JoyImageEditPipeline) + pipeline.vae_scale_factor = 8 + pipeline.latent_channels = 4 + image_latents = torch.ones(1, 4, 1, 2, 2) + noise_latents = torch.zeros(1, 1, 4, 1, 2, 2) + pipeline._encode_vae_image = lambda image, generator: image_latents + + latents, reference_latents = JoyImageEditPipeline._prepare_latents( + pipeline, + image=torch.zeros(1, 3, 1, 16, 16), + batch_size=1, + num_channels_latents=4, + height=16, + width=16, + dtype=torch.float32, + device=torch.device("cpu"), + generator=None, + latents=noise_latents, + ) + + assert latents.shape == (1, 2, 4, 1, 2, 2) + assert torch.equal(latents[:, :1], reference_latents) + assert torch.equal(latents[:, -1:], noise_latents) + + +def test_prepare_latents_validates_generator_and_latent_shapes(): + pipeline = object.__new__(JoyImageEditPipeline) + pipeline.vae_scale_factor = 8 + pipeline.latent_channels = 4 + image_latents = torch.ones(1, 4, 1, 2, 2) + pipeline._encode_vae_image = lambda image, generator: image_latents + + with pytest.raises(ValueError, match="Generator list length"): + JoyImageEditPipeline._prepare_latents( + pipeline, + image=torch.zeros(1, 3, 1, 16, 16), + batch_size=2, + num_channels_latents=4, + height=16, + width=16, + dtype=torch.float32, + device=torch.device("cpu"), + generator=[torch.Generator()], + ) + + with pytest.raises(ValueError, match="noise latents must have shape"): + JoyImageEditPipeline._prepare_latents( + pipeline, + image=torch.zeros(1, 3, 1, 16, 16), + batch_size=1, + num_channels_latents=4, + height=16, + width=16, + dtype=torch.float32, + device=torch.device("cpu"), + generator=None, + latents=torch.zeros(1, 1, 4, 1, 1, 2), + ) + + +def test_prepare_latents_validates_image_latent_shape(): + pipeline = object.__new__(JoyImageEditPipeline) + pipeline.vae_scale_factor = 8 + pipeline.latent_channels = 4 + pipeline._encode_vae_image = lambda image, generator: torch.ones(1, 4, 2, 2) + + with pytest.raises(ValueError, match="image latents must have shape"): + JoyImageEditPipeline._prepare_latents( + pipeline, + image=torch.zeros(1, 3, 1, 16, 16), + batch_size=1, + num_channels_latents=4, + height=16, + width=16, + dtype=torch.float32, + device=torch.device("cpu"), + generator=None, + ) + + +def test_diffuse_restores_reference_slots_each_step(): + class FakeScheduler: + def set_begin_index(self, index): + self.begin_index = index + + def step(self, noise_pred, timestep, latents, return_dict=False): + return (latents + 1.0,) + + class FakeTransformer(torch.nn.Module): + def forward(self, hidden_states, **kwargs): + return (torch.zeros_like(hidden_states),) + + pipeline = object.__new__(JoyImageEditPipeline) + torch.nn.Module.__init__(pipeline) + pipeline.scheduler = FakeScheduler() + pipeline.transformer = FakeTransformer() + pipeline._interrupt = False + pipeline._current_timestep = None + + latents = torch.zeros(1, 2, 4, 1, 2, 2) + image_latents = torch.full((1, 1, 4, 1, 2, 2), 7.0) + result = JoyImageEditPipeline.diffuse( + pipeline, + latents=latents, + image_latents=image_latents, + prompt_embeds=torch.zeros(1, 3, 8), + prompt_embeds_mask=torch.ones(1, 3, dtype=torch.long), + negative_prompt_embeds=None, + negative_prompt_embeds_mask=None, + timesteps=torch.tensor([2.0, 1.0]), + do_true_cfg=False, + true_cfg_scale=1.0, + ) + + assert torch.equal(result[:, :1], image_latents) + assert torch.equal(result[:, -1:], torch.full((1, 1, 4, 1, 2, 2), 2.0)) + + +@pytest.mark.parametrize( + "params_kwargs, expected_scale", + [ + ({}, 4.0), + ({"guidance_scale": 2.0}, 2.0), + ({"guidance_scale": 1.0}, 1.0), + ({"guidance_scale": 0.0}, 0.0), + ({"true_cfg_scale": 8.0}, 8.0), + ], +) +def test_forward_uses_guidance_scale_for_cfg(params_kwargs, expected_scale): + class FakeScheduler: + def set_timesteps(self, num_inference_steps, device): + self.num_inference_steps = num_inference_steps + self.timesteps = torch.tensor([1.0], device=device) + + pipeline = object.__new__(JoyImageEditPipeline) + torch.nn.Module.__init__(pipeline) + pipeline.device = torch.device("cpu") + pipeline.latent_channels = 4 + pipeline.scheduler = FakeScheduler() + encode_calls = [] + diffuse_calls = [] + + def encode_prompt(prompt, image, **kwargs): + encode_calls.append(prompt) + return torch.zeros(1, 3, 4), torch.ones(1, 3, dtype=torch.long) + + def fake_prepare_latents(**kwargs): + latents = torch.zeros(1, 2, 4, 1, 2, 2) + image_latents = torch.ones(1, 1, 4, 1, 2, 2) + return latents, image_latents + + def diffuse(**kwargs): + diffuse_calls.append(kwargs) + return kwargs["latents"] + + pipeline.encode_prompt = encode_prompt + pipeline._prepare_latents = fake_prepare_latents + pipeline.diffuse = diffuse + pipeline._decode_latents = lambda latents: torch.zeros(1, 3, 2, 2) + + request = _make_request( + prompt={ + "prompt": "make it brighter", + "additional_information": { + "image_tensor": torch.zeros(1, 3, 1, 16, 16), + "prompt_image": Image.new("RGB", (16, 16)), + "height": 16, + "width": 16, + }, + }, + params=_make_params(**params_kwargs), + ) + + output = JoyImageEditPipeline.forward(pipeline, request) + + assert output.output.shape == (1, 3, 2, 2) + expected_cfg = expected_scale > 1.0 + assert encode_calls == (["make it brighter", ""] if expected_cfg else ["make it brighter"]) + assert diffuse_calls[0]["do_true_cfg"] is expected_cfg + assert diffuse_calls[0]["true_cfg_scale"] == expected_scale + + +@pytest.mark.parametrize( + "schedule_kwargs, expected_timesteps", + [ + ({"timesteps": torch.tensor([9, 4])}, torch.tensor([9.0, 4.0])), + ({"sigmas": [1.0, 0.5]}, torch.tensor([1000.0, 500.0])), + ], +) +def test_forward_uses_request_sampling_parameters(schedule_kwargs, expected_timesteps): + class FakeScheduler: + def set_timesteps(self, num_inference_steps=None, device=None, timesteps=None, sigmas=None): + if timesteps is not None: + self.timesteps = torch.as_tensor(timesteps, device=device, dtype=torch.float32) + elif sigmas is not None: + self.timesteps = torch.as_tensor(sigmas, device=device) * 1000 + else: + self.timesteps = torch.arange(num_inference_steps, device=device, dtype=torch.float32) + + pipeline = object.__new__(JoyImageEditPipeline) + torch.nn.Module.__init__(pipeline) + pipeline.device = torch.device("cpu") + pipeline.latent_channels = 4 + pipeline.tokenizer_max_length = JOY_MAX_IMAGE_SEQ_LEN + pipeline.scheduler = FakeScheduler() + encode_calls = [] + prepare_calls = [] + diffuse_calls = [] + + def encode_prompt(prompt, image, **kwargs): + encode_calls.append(kwargs) + return torch.zeros(1, 3, 4), torch.ones(1, 3, dtype=torch.long) + + def prepare_latents(**kwargs): + prepare_calls.append(kwargs) + reference = torch.ones(1, 1, 4, 1, 2, 2) + return torch.cat([reference, kwargs["latents"]], dim=1), reference + + def diffuse(**kwargs): + diffuse_calls.append(kwargs) + return kwargs["latents"] + + pipeline.encode_prompt = encode_prompt + pipeline._prepare_latents = prepare_latents + pipeline.diffuse = diffuse + pipeline._decode_latents = lambda _latents: pytest.fail("latent output must not be decoded") + + request_latents = torch.full((1, 1, 4, 1, 2, 2), 3.0) + params = _make_params( + guidance_scale=1.0, + latents=request_latents, + output_type="latent", + max_sequence_length=123, + **schedule_kwargs, + ) + request = _make_request( + prompt={ + "prompt": "make it brighter", + "additional_information": { + "image_tensor": torch.zeros(1, 3, 1, 16, 16), + "prompt_image": Image.new("RGB", (16, 16)), + "height": 16, + "width": 16, + }, + }, + params=params, + ) + + output = JoyImageEditPipeline.forward(pipeline, request) + + assert encode_calls[0]["max_sequence_length"] == 123 + assert prepare_calls[0]["latents"] is request_latents + torch.testing.assert_close(diffuse_calls[0]["timesteps"], expected_timesteps) + assert diffuse_calls[0]["cfg_normalize"] is True + torch.testing.assert_close(output.output, request_latents[:, 0]) + + +def test_retrieve_timesteps_rejects_conflicting_custom_schedules(): + scheduler = SimpleNamespace(set_timesteps=lambda **_kwargs: None) + + with pytest.raises(ValueError, match="Only one"): + retrieve_timesteps(scheduler, timesteps=[9], sigmas=[1.0]) + + +def test_forward_rejects_unsupported_output_type(): + pipeline = object.__new__(JoyImageEditPipeline) + torch.nn.Module.__init__(pipeline) + request = _make_request( + prompt={ + "prompt": "make it brighter", + "additional_information": { + "image_tensor": torch.zeros(1, 3, 1, 16, 16), + "prompt_image": Image.new("RGB", (16, 16)), + "height": 16, + "width": 16, + }, + }, + params=_make_params(output_type="np"), + ) + + with pytest.raises(ValueError, match="only `pil` and `latent`"): + JoyImageEditPipeline.forward(pipeline, request) + + +def test_forward_skips_decode_for_dummy_warmup_request(monkeypatch): + class FakeScheduler: + def set_timesteps(self, num_inference_steps, device): + self.timesteps = torch.tensor([1.0], device=device) + + pipeline = object.__new__(JoyImageEditPipeline) + torch.nn.Module.__init__(pipeline) + pipeline.device = torch.device("cpu") + pipeline.od_config = SimpleNamespace( + enable_cpu_offload=True, + enable_layerwise_offload=False, + pin_cpu_memory=False, + parallel_config=SimpleNamespace(use_hsdp=False), + ) + pipeline.transformer = torch.nn.Linear(1, 1) + pipeline.latent_channels = 4 + pipeline.scheduler = FakeScheduler() + pipeline.encode_prompt = lambda *args, **kwargs: ( + torch.zeros(1, 3, 4), + torch.ones(1, 3, dtype=torch.long), + ) + pipeline._prepare_latents = lambda **kwargs: ( + torch.zeros(1, 2, 4, 1, 2, 2), + torch.ones(1, 1, 4, 1, 2, 2), + ) + pipeline.diffuse = lambda **kwargs: kwargs["latents"] + 1.0 + pipeline.check_cfg_parallel_validity = lambda scale, has_neg_prompt: True + + def decode_should_not_run(latents): + raise AssertionError("dummy warmup should not decode latents") + + pipeline._decode_latents = decode_should_not_run + move_calls = [] + + def fake_move_params(module, target_device, **kwargs): + move_calls.append((module, target_device, kwargs)) + + monkeypatch.setattr( + joy_pipeline_module.SequentialOffloadHook, + "_move_params", + fake_move_params, + ) + monkeypatch.setattr(joy_pipeline_module.current_omni_platform, "empty_cache", lambda: None) + + request = _make_request( + prompt={ + "prompt": "dummy run", + "additional_information": { + "image_tensor": torch.zeros(1, 3, 1, 16, 16), + "prompt_image": Image.new("RGB", (16, 16)), + "height": 16, + "width": 16, + }, + }, + params=_make_params(guidance_scale=1.0), + ) + request.request_id = DUMMY_DIFFUSION_REQUEST_ID + request.is_dummy_run = lambda: True + + output = JoyImageEditPipeline.forward(pipeline, request) + + assert output.output.shape == (1, 4, 1, 2, 2) + assert move_calls == [ + ( + pipeline.transformer, + torch.device("cpu"), + {"non_blocking": False, "pin_memory": False}, + ) + ] diff --git a/tests/e2e/accuracy/test_joy_image_edit.py b/tests/e2e/accuracy/test_joy_image_edit.py new file mode 100644 index 000000000..82d0ca411 --- /dev/null +++ b/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, + ) + ) diff --git a/tests/model_tests/diffusion/test_alignment.py b/tests/model_tests/diffusion/test_alignment.py index 5daeded11..b5345eb87 100644 --- a/tests/model_tests/diffusion/test_alignment.py +++ b/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", diff --git a/vllm_omni/diffusion/model_metadata.py b/vllm_omni/diffusion/model_metadata.py index ba91e2e21..261341f5d 100644 --- a/vllm_omni/diffusion/model_metadata.py +++ b/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, diff --git a/vllm_omni/diffusion/models/joy_image/__init__.py b/vllm_omni/diffusion/models/joy_image/__init__.py new file mode 100644 index 000000000..e50b6d4d7 --- /dev/null +++ b/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", +] diff --git a/vllm_omni/diffusion/models/joy_image/cfg_parallel.py b/vllm_omni/diffusion/models/joy_image/cfg_parallel.py new file mode 100644 index 000000000..9fd235d63 --- /dev/null +++ b/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 diff --git a/vllm_omni/diffusion/models/joy_image/joy_image_edit_transformer.py b/vllm_omni/diffusion/models/joy_image/joy_image_edit_transformer.py new file mode 100644 index 000000000..248db623b --- /dev/null +++ b/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 diff --git a/vllm_omni/diffusion/models/joy_image/pipeline_joy_image_edit.py b/vllm_omni/diffusion/models/joy_image/pipeline_joy_image_edit.py new file mode 100644 index 000000000..adf945d3d --- /dev/null +++ b/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"\n{prompt}" + prompt = f"<|im_start|>user\n{prompt}<|im_end|>\n" + return prompt.replace("\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) diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index 2c4cfb024..690ced068 100644 --- a/vllm_omni/diffusion/registry.py +++ b/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",