diff --git a/.github/workflows/build_documentation.yml b/.github/workflows/build_documentation.yml index 5bf7fe5daf5c..3bcf7cc77345 100644 --- a/.github/workflows/build_documentation.yml +++ b/.github/workflows/build_documentation.yml @@ -17,14 +17,13 @@ permissions: jobs: build: - uses: huggingface/doc-builder/.github/workflows/build_main_documentation.yml@2430c1ec91d04667414e2fa31ecfc36c153ea391 # main + uses: huggingface/doc-builder/.github/workflows/build_main_documentation.yml@e60a538eea9817ab312196d0d233604b01697265 # main with: commit_sha: ${{ github.sha }} install_libgl1: true package: diffusers notebook_folder: diffusers_doc languages: en ko zh ja pt - custom_container: diffusers/diffusers-doc-builder pre_command: uv pip uninstall transformers huggingface_hub && UV_PRERELEASE=allow uv pip install -U transformers@git+https://github.com/huggingface/transformers.git secrets: token: ${{ secrets.HUGGINGFACE_PUSH }} diff --git a/.github/workflows/build_pr_documentation.yml b/.github/workflows/build_pr_documentation.yml index 8bc015cdecf2..5b29114cf251 100644 --- a/.github/workflows/build_pr_documentation.yml +++ b/.github/workflows/build_pr_documentation.yml @@ -42,12 +42,11 @@ jobs: build: needs: check-links - uses: huggingface/doc-builder/.github/workflows/build_pr_documentation.yml@90b4ee2c10b81b5c1a6367c4e6fc9e2fb510a7e3 # main + uses: huggingface/doc-builder/.github/workflows/build_pr_documentation.yml@e60a538eea9817ab312196d0d233604b01697265 # main with: commit_sha: ${{ github.event.pull_request.head.sha }} pr_number: ${{ github.event.number }} install_libgl1: true package: diffusers languages: en ko zh ja pt - custom_container: diffusers/diffusers-doc-builder pre_command: uv pip uninstall transformers huggingface_hub && UV_PRERELEASE=allow uv pip install -U transformers@git+https://github.com/huggingface/transformers.git diff --git a/docker/diffusers-pytorch-cuda/Dockerfile b/docker/diffusers-pytorch-cuda/Dockerfile index b6eedaad50ba..28ae25717c54 100644 --- a/docker/diffusers-pytorch-cuda/Dockerfile +++ b/docker/diffusers-pytorch-cuda/Dockerfile @@ -40,11 +40,6 @@ RUN uv pip install --no-cache-dir \ torchaudio==2.10.0 \ --index-url https://download.pytorch.org/whl/cu129 -# Install compatible versions of numba/llvmlite for Python 3.10+ -RUN uv pip install --no-cache-dir \ - "llvmlite>=0.40.0" \ - "numba>=0.57.0" - RUN uv pip install --no-cache-dir "git+https://github.com/huggingface/diffusers.git@main#egg=diffusers[test]" # Extra dependencies diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index 83bc19670d96..0ae5262c39e3 100644 --- a/docs/source/en/_toctree.yml +++ b/docs/source/en/_toctree.yml @@ -353,6 +353,8 @@ title: HunyuanVideoTransformer3DModel - local: api/models/ideogram4_transformer2d title: Ideogram4Transformer2DModel + - local: api/models/transformer_joyimage_edit_plus + title: JoyImageEditPlusTransformer3DModel - local: api/models/transformer_joyimage title: JoyImageEditTransformer3DModel - local: api/models/krea2_transformer2d @@ -555,6 +557,8 @@ title: InstructPix2Pix - local: api/pipelines/joyimage_edit title: JoyImage Edit + - local: api/pipelines/joyimage_edit_plus + title: JoyImage Edit Plus - local: api/pipelines/kandinsky title: Kandinsky 2.1 - local: api/pipelines/kandinsky_v22 diff --git a/docs/source/en/api/models/transformer_joyimage_edit_plus.md b/docs/source/en/api/models/transformer_joyimage_edit_plus.md new file mode 100644 index 000000000000..776c53eaf20c --- /dev/null +++ b/docs/source/en/api/models/transformer_joyimage_edit_plus.md @@ -0,0 +1,29 @@ + + +# JoyImageEditPlusTransformer3DModel + +The model can be loaded with the following code snippet. + +```python +from diffusers import JoyImageEditPlusTransformer3DModel + +transformer = JoyImageEditPlusTransformer3DModel.from_pretrained("jdopensource/JoyAI-Image-Edit-Plus-Diffusers", subfolder="transformer", torch_dtype=torch.bfloat16) +``` + +## JoyImageEditPlusTransformer3DModel + +[[autodoc]] JoyImageEditPlusTransformer3DModel + +## Transformer2DModelOutput + +[[autodoc]] models.modeling_outputs.Transformer2DModelOutput diff --git a/docs/source/en/api/pipelines/animatediff.md b/docs/source/en/api/pipelines/animatediff.md index f0188f3c36fb..60057e15eedc 100644 --- a/docs/source/en/api/pipelines/animatediff.md +++ b/docs/source/en/api/pipelines/animatediff.md @@ -68,7 +68,7 @@ scheduler = DDIMScheduler.from_pretrained( pipe.scheduler = scheduler # enable memory savings -pipe.enable_vae_slicing() +pipe.vae.enable_slicing() pipe.enable_model_cpu_offload() output = pipe( @@ -395,8 +395,8 @@ pipe = AnimateDiffSDXLPipeline.from_pretrained( ).to("cuda") # enable memory savings -pipe.enable_vae_slicing() -pipe.enable_vae_tiling() +pipe.vae.enable_slicing() +pipe.vae.enable_tiling() output = pipe( prompt="a panda surfing in the ocean, realistic, high quality", @@ -441,7 +441,7 @@ scheduler = DDIMScheduler.from_pretrained( pipe.scheduler = scheduler # enable memory savings -pipe.enable_vae_slicing() +pipe.vae.enable_slicing() pipe.enable_model_cpu_offload() # helper function to load videos @@ -640,7 +640,7 @@ scheduler = DDIMScheduler.from_pretrained( pipe.scheduler = scheduler # enable memory savings -pipe.enable_vae_slicing() +pipe.vae.enable_slicing() pipe.enable_model_cpu_offload() output = pipe( @@ -714,7 +714,7 @@ scheduler = DDIMScheduler.from_pretrained( pipe.scheduler = scheduler # enable memory savings -pipe.enable_vae_slicing() +pipe.vae.enable_slicing() pipe.enable_model_cpu_offload() output = pipe( @@ -772,8 +772,8 @@ pipe.scheduler = DDIMScheduler.from_pretrained( ) # enable memory savings -pipe.enable_vae_slicing() -pipe.enable_vae_tiling() +pipe.vae.enable_slicing() +pipe.vae.enable_tiling() # enable FreeInit # Refer to the enable_free_init documentation for a full list of configurable parameters @@ -840,7 +840,7 @@ pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config, beta_schedule=" pipe.load_lora_weights("wangfuyun/AnimateLCM", weight_name="sd15_lora_beta.safetensors", adapter_name="lcm-lora") -pipe.enable_vae_slicing() +pipe.vae.enable_slicing() pipe.enable_model_cpu_offload() output = pipe( @@ -882,7 +882,7 @@ pipe.load_lora_weights("wangfuyun/AnimateLCM", weight_name="sd15_lora_beta.safet pipe.load_lora_weights("guoyww/animatediff-motion-lora-tilt-up", adapter_name="tilt-up") pipe.set_adapters(["lcm-lora", "tilt-up"], [1.0, 0.8]) -pipe.enable_vae_slicing() +pipe.vae.enable_slicing() pipe.enable_model_cpu_offload() output = pipe( diff --git a/docs/source/en/api/pipelines/controlnet.md b/docs/source/en/api/pipelines/controlnet.md index afc0a4653e07..9e7d387e4e7c 100644 --- a/docs/source/en/api/pipelines/controlnet.md +++ b/docs/source/en/api/pipelines/controlnet.md @@ -37,8 +37,6 @@ The original codebase can be found at [lllyasviel/ControlNet](https://github.com - __call__ - enable_attention_slicing - disable_attention_slicing - - enable_vae_slicing - - disable_vae_slicing - enable_xformers_memory_efficient_attention - disable_xformers_memory_efficient_attention - load_textual_inversion @@ -49,8 +47,6 @@ The original codebase can be found at [lllyasviel/ControlNet](https://github.com - __call__ - enable_attention_slicing - disable_attention_slicing - - enable_vae_slicing - - disable_vae_slicing - enable_xformers_memory_efficient_attention - disable_xformers_memory_efficient_attention - load_textual_inversion @@ -61,8 +57,6 @@ The original codebase can be found at [lllyasviel/ControlNet](https://github.com - __call__ - enable_attention_slicing - disable_attention_slicing - - enable_vae_slicing - - disable_vae_slicing - enable_xformers_memory_efficient_attention - disable_xformers_memory_efficient_attention - load_textual_inversion diff --git a/docs/source/en/api/pipelines/cosmos3.md b/docs/source/en/api/pipelines/cosmos3.md index 1ac8f36457a4..93efbbb5e51e 100644 --- a/docs/source/en/api/pipelines/cosmos3.md +++ b/docs/source/en/api/pipelines/cosmos3.md @@ -738,10 +738,266 @@ pipe = Cosmos3OmniPipeline.from_pretrained( - all - __call__ +## Cosmos3OmniModularPipeline + +Cosmos 3 is also available as a Modular Diffusers pipeline. The task-based [`Cosmos3OmniPipeline`] remains available; the modular pipeline coexists with it and covers the same modes (`text2image`, `text2video`, `image2video`, `video2video`, action-conditioned generation, and `transfer` (structural control), with optional sound when supported by the checkpoint). + +```python +import torch +from diffusers import Cosmos3OmniModularPipeline + +pipe = Cosmos3OmniModularPipeline.from_pretrained( + "nvidia/Cosmos3-Nano", torch_dtype=torch.bfloat16 +) +pipe.load_components(torch_dtype=torch.bfloat16) +pipe.enable_safety_checker() + +videos = pipe( + prompt='{"scene":"A robot arm in a kitchen"}', + num_frames=1, + height=720, + width=1280, + output="videos", +) + +# Modular pipelines expose declared outputs directly instead of using the task pipeline's +# `return_dict`/`Cosmos3OmniPipelineOutput` API. +image = videos[0] +``` + +You can also load through [`ModularPipeline`] and let the repository config select the blocks class: + +```python +import torch +from diffusers import ModularPipeline + +pipe = ModularPipeline.from_pretrained("nvidia/Cosmos3-Nano", torch_dtype=torch.bfloat16) +pipe.load_components(torch_dtype=torch.bfloat16) +pipe.enable_safety_checker() +videos = pipe( + prompt='{"scene":"A robot arm in a kitchen"}', num_frames=1, height=720, width=1280, output="videos" +) +``` + +To inspect or customize a specific Cosmos modular workflow, use `available_workflows` + `get_workflow()`: + +```python +available = pipe.blocks.available_workflows +image2video_blocks = pipe.blocks.get_workflow("image2video") +``` + +### Modular examples for all existing workflows + +The modular pipeline supports the same call signatures as the task pipeline. The snippets below mirror every generation example shown above (`text2video`, `text2image`, `image2video`, `video2video`, `video2video_sound`, `text2video_sound`, and `action_policy`). Transfer (structural control) has its own inputs and is shown separately in [Modular transfer](#modular-transfer-structural-control) below. + +```python +import json +import torch +from diffusers import Cosmos3OmniModularPipeline, CosmosActionCondition +from diffusers.schedulers.scheduling_unipc_multistep import UniPCMultistepScheduler +from diffusers.utils import encode_video, export_to_video, load_image, load_video + +pipe = Cosmos3OmniModularPipeline.from_pretrained("nvidia/Cosmos3-Nano", torch_dtype=torch.bfloat16) +pipe.load_components(torch_dtype=torch.bfloat16) +pipe.enable_safety_checker() +pipe.to("cuda") +pipe.scheduler = UniPCMultistepScheduler.from_config( + pipe.scheduler.config, flow_shift=10.0, use_karras_sigmas=False +) + +# text2video +json_prompt = json.load(open("assets/example_t2v_prompt.json")) +negative_prompt = json.load(open("assets/negative_prompt.json")) +videos = pipe( + prompt=json.dumps(json_prompt), + negative_prompt=json.dumps(negative_prompt), + num_frames=189, + height=720, + width=1280, + num_inference_steps=35, + guidance_scale=6.0, + fps=24.0, + output="videos", +) +export_to_video(videos, "cosmos3_modular_t2v.mp4", fps=24, macro_block_size=1) + +# text2image +json_prompt = json.load(open("assets/example_t2i_prompt.json")) +videos = pipe(prompt=json.dumps(json_prompt), num_frames=1, height=720, width=1280, output="videos") +videos[0].save("cosmos3_modular_t2i.jpg", format="JPEG", quality=85) + +# image2video +json_prompt = json.load(open("assets/example_i2v_prompt.json")) +negative_prompt = json.load(open("assets/negative_prompt_i2v.json")) +image = load_image("https://github.com/nvidia-cosmos/cosmos-dependencies/releases/download/assets/robot_153.jpg") +videos = pipe( + prompt=json.dumps(json_prompt), + negative_prompt=json.dumps(negative_prompt), + image=image, + num_frames=189, + height=720, + width=1280, + fps=24.0, + output="videos", +) +export_to_video(videos, "cosmos3_modular_i2v.mp4", fps=24, macro_block_size=1) + +# video2video +json_prompt = json.load(open("assets/example_v2v_prompt.json")) +negative_prompt = json.load(open("assets/negative_prompt_i2v.json")) +video = load_video( + "https://github.com/nvidia-cosmos/cosmos-dependencies/raw/refs/heads/assets/cosmos3/inputs/vision/robot_pouring.mp4" +) +videos = pipe( + prompt=json.dumps(json_prompt), + negative_prompt=json.dumps(negative_prompt), + video=video, + condition_frame_indexes_vision=[0, 1], + condition_video_keep="first", + num_frames=189, + height=720, + width=1280, + num_inference_steps=35, + guidance_scale=6.0, + fps=24.0, + output="videos", +) +export_to_video(videos, "cosmos3_modular_v2v.mp4", fps=24, macro_block_size=1) + +# video2video_sound +outputs = pipe( + prompt=json.dumps(json_prompt), + negative_prompt=json.dumps(negative_prompt), + video=video, + condition_frame_indexes_vision=[0, 1], + condition_video_keep="first", + num_frames=189, + height=720, + width=1280, + fps=24.0, + enable_sound=True, + output=["videos", "sound", "sampling_rate"], +) +encode_video( + outputs["videos"], + fps=24, + audio=outputs["sound"], + audio_sample_rate=outputs["sampling_rate"], + output_path="cosmos3_modular_v2v_with_sound.mp4", +) + +# text2video_sound +json_prompt = json.load(open("assets/example_t2v_sound_prompt.json")) +negative_prompt = json.load(open("assets/negative_prompt.json")) +outputs = pipe( + prompt=json.dumps(json_prompt), + negative_prompt=json.dumps(negative_prompt), + num_frames=189, + height=720, + width=1280, + fps=24.0, + enable_sound=True, + output=["videos", "sound", "sampling_rate"], +) +encode_video( + outputs["videos"], + fps=24, + audio=outputs["sound"], + audio_sample_rate=outputs["sampling_rate"], + output_path="cosmos3_modular_t2v_with_sound.mp4", +) + +# action_policy +prompt = "Put the pot to the left of the purple item." +action_video = load_video( + "https://github.com/nvidia-cosmos/cosmos-dependencies/raw/refs/heads/assets/cosmos3/inputs/action/bridge_20260501_0.mp4" +) +outputs = pipe( + prompt=prompt, + action=CosmosActionCondition( + mode="policy", + chunk_size=16, + domain_name="bridge_orig_lerobot", + resolution_tier=480, + video=action_video, + view_point="ego_view", + ), + fps=5, + num_inference_steps=30, + guidance_scale=1.0, + use_system_prompt=False, + output=["videos", "action"], +) +export_to_video(outputs["videos"], "cosmos3_modular_action_policy.mp4", fps=5, macro_block_size=1) +if outputs["action"] is not None: + with open("cosmos3_modular_action_policy.json", "w") as f: + json.dump(outputs["action"][0].tolist(), f) +``` + +### Modular transfer (structural control) + +Transfer follows a **precomputed control video** (edge, blur, depth, segmentation, or a world-scenario map) passed through `control_videos=` as a `{hint: video}` mapping. It is video-only (no `image` / `video` / `action` / `enable_sound`), the prompt is a pre-upsampled JSON caption (see [Prompt upsampling](#prompt-upsampling)), and long clips are generated autoregressively in chunks of `num_video_frames_per_chunk` and stitched automatically. `guidance_scale` is the usual text CFG; `control_guidance` (`!= 1.0`) additionally amplifies the control signal. Recommended starting values per hint: + +| Hint | `guidance_scale` | `control_guidance` | `flow_shift` | Geometry | +| --- | --- | --- | --- | --- | +| Edge / Blur / Depth | 3.0 | 1.5 | 10.0 | 121 frames @ 30 FPS | +| Segmentation | 3.0 | 2.0 | 10.0 | 121 frames @ 30 FPS | +| World scenario (WSM) | 1.0 | 3.0 | 10.0 | 101 frames @ 10 FPS | + +Diffusers does not ship the control assets. Ready-made ones (a control video + matching `prompt.json` per hint, plus a shared `negative_prompt.json`) live in the [Cosmos cookbook](https://github.com/NVIDIA/cosmos/tree/main/cookbooks/cosmos3/generator/transfer/assets). For the edge example below, download them into a local `assets/` folder: + +```bash +base=https://github.com/NVIDIA/cosmos/raw/refs/heads/main/cookbooks/cosmos3/generator/transfer/assets +mkdir -p assets/edge +curl -sL "$base/edge/control_edge.mp4" -o assets/edge/control_edge.mp4 +curl -sL "$base/edge/prompt.json" -o assets/edge/prompt.json +curl -sL "$base/negative_prompt.json" -o assets/negative_prompt.json +``` + +```python +import json +import torch +from diffusers import Cosmos3OmniModularPipeline +from diffusers.schedulers.scheduling_unipc_multistep import UniPCMultistepScheduler +from diffusers.utils import export_to_video, load_video + +pipe = Cosmos3OmniModularPipeline.from_pretrained("nvidia/Cosmos3-Nano", torch_dtype=torch.bfloat16) +pipe.load_components(torch_dtype=torch.bfloat16) +pipe.to("cuda") +pipe.scheduler = UniPCMultistepScheduler.from_config( + pipe.scheduler.config, flow_shift=10.0, use_karras_sigmas=False +) + +# Downloaded into assets/ from the Cosmos cookbook (see the curl snippet above). +json_prompt = json.load(open("assets/edge/prompt.json")) +negative_prompt = json.load(open("assets/negative_prompt.json")) +control_edge = load_video("assets/edge/control_edge.mp4") + +videos = pipe( + prompt=json.dumps(json_prompt), + negative_prompt=json.dumps(negative_prompt), + control_videos={"edge": control_edge}, + num_frames=121, + height=720, + width=1280, + fps=30.0, + num_inference_steps=35, + guidance_scale=3.0, + control_guidance=1.5, + output="videos", +) +export_to_video(videos, "cosmos3_modular_transfer_edge.mp4", fps=30, macro_block_size=1) +``` + +[[autodoc]] Cosmos3OmniModularPipeline + +- all +- __call__ + ## CosmosActionCondition [[autodoc]] CosmosActionCondition ## Cosmos3OmniPipelineOutput -[[autodoc]] pipelines.cosmos.pipeline_cosmos3_omni.Cosmos3OmniPipelineOutput \ No newline at end of file +[[autodoc]] pipelines.cosmos.pipeline_cosmos3_omni.Cosmos3OmniPipelineOutput diff --git a/docs/source/en/api/pipelines/joyimage_edit_plus.md b/docs/source/en/api/pipelines/joyimage_edit_plus.md new file mode 100644 index 000000000000..92609d166af0 --- /dev/null +++ b/docs/source/en/api/pipelines/joyimage_edit_plus.md @@ -0,0 +1,61 @@ + + +# JoyAI-Image-Edit-Plus + +[JoyAI-Image](https://github.com/jd-opensource/JoyAI-Image) is a unified multimodal foundation model for image understanding, text-to-image generation, and instruction-guided image editing. It combines an 8B Multimodal Large Language Model (MLLM) with a 16B Multimodal Diffusion Transformer (MMDiT). + +JoyAI-Image-Edit-Plus is a multi-image instruction-guided editing model that accepts **multiple reference images** and a text instruction to generate a new image that combines elements from the references according to the instruction. It supports 1–5 reference images per sample. + +| Model | Description | Download | +|:-----:|:-----------:|:--------:| +| JoyAI-Image-Edit-Plus | Multi-image instruction-guided editing with element composition from multiple references | [Hugging Face](https://huggingface.co/jdopensource/JoyAI-Image-Edit-Plus-Diffusers) | + +```python +import torch +from PIL import Image +from diffusers import JoyImageEditPlusPipeline + +pipeline = JoyImageEditPlusPipeline.from_pretrained( + "jdopensource/JoyAI-Image-Edit-Plus-Diffusers", torch_dtype=torch.bfloat16 +) +pipeline.to("cuda") + +images = [ + Image.open("reference_0.png").convert("RGB"), + Image.open("reference_1.png").convert("RGB"), +] + +target_h, target_w = pipeline.image_processor.get_default_height_width(images[-1]) + +output = pipeline( + images=images, + prompt="Combine the person from the second image with the scene from the first image.", + negative_prompt="low quality, blurry, deformed", + height=target_h, + width=target_w, + num_inference_steps=30, + guidance_scale=4.0, + generator=torch.Generator("cuda").manual_seed(42), +).images[0] +output.save("joyimage_edit_plus_output.png") +``` + +## JoyImageEditPlusPipeline + +[[autodoc]] JoyImageEditPlusPipeline + - all + - __call__ + +## JoyImageEditPlusPipelineOutput + +[[autodoc]] pipelines.joyimage.pipeline_output.JoyImageEditPlusPipelineOutput diff --git a/docs/source/en/api/pipelines/latent_consistency_models.md b/docs/source/en/api/pipelines/latent_consistency_models.md index aee6dda64fa1..14aee3de2eae 100644 --- a/docs/source/en/api/pipelines/latent_consistency_models.md +++ b/docs/source/en/api/pipelines/latent_consistency_models.md @@ -643,10 +643,6 @@ export_to_gif(frames, "animation.gif") - __call__ - enable_freeu - disable_freeu - - enable_vae_slicing - - disable_vae_slicing - - enable_vae_tiling - - disable_vae_tiling ## LatentConsistencyModelImg2ImgPipeline @@ -655,10 +651,6 @@ export_to_gif(frames, "animation.gif") - __call__ - enable_freeu - disable_freeu - - enable_vae_slicing - - disable_vae_slicing - - enable_vae_tiling - - disable_vae_tiling ## StableDiffusionPipelineOutput diff --git a/docs/source/en/api/pipelines/mochi.md b/docs/source/en/api/pipelines/mochi.md index f19a9bd575c1..f9daa3515774 100644 --- a/docs/source/en/api/pipelines/mochi.md +++ b/docs/source/en/api/pipelines/mochi.md @@ -86,7 +86,7 @@ pipe = MochiPipeline.from_pretrained("genmo/mochi-1-preview") # Enable memory savings pipe.enable_model_cpu_offload() -pipe.enable_vae_tiling() +pipe.vae.enable_tiling() prompt = "Close-up of a chameleon's eye, with its scaly skin changing color. Ultra high resolution 4k." @@ -109,7 +109,7 @@ pipe = MochiPipeline.from_pretrained("genmo/mochi-1-preview", variant="bf16", to # Enable memory savings pipe.enable_model_cpu_offload() -pipe.enable_vae_tiling() +pipe.vae.enable_tiling() prompt = "Close-up of a chameleon's eye, with its scaly skin changing color. Ultra high resolution 4k." frames = pipe(prompt, num_frames=85).frames[0] @@ -138,7 +138,7 @@ from diffusers.utils import export_to_video from diffusers.video_processor import VideoProcessor pipe = MochiPipeline.from_pretrained("genmo/mochi-1-preview", force_zeros_for_empty_prompt=True) -pipe.enable_vae_tiling() +pipe.vae.enable_tiling() pipe.enable_model_cpu_offload() prompt = "An aerial shot of a parade of elephants walking across the African savannah. The camera showcases the herd and the surrounding landscape." @@ -205,7 +205,7 @@ transformer = MochiTransformer3DModel.from_pretrained( pipe = MochiPipeline.from_pretrained(model_id, transformer=transformer) pipe.enable_model_cpu_offload() -pipe.enable_vae_tiling() +pipe.vae.enable_tiling() with torch.autocast(device_type="cuda", dtype=torch.bfloat16, cache_enabled=False): frames = pipe( @@ -245,7 +245,7 @@ transformer = MochiTransformer3DModel.from_pretrained(ckpt_path, torch_dtype=tor pipe = MochiPipeline.from_pretrained(model_id, transformer=transformer) pipe.enable_model_cpu_offload() -pipe.enable_vae_tiling() +pipe.vae.enable_tiling() with torch.autocast(device_type="cuda", dtype=torch.bfloat16, cache_enabled=False): frames = pipe( diff --git a/docs/source/en/api/pipelines/stable_diffusion/adapter.md b/docs/source/en/api/pipelines/stable_diffusion/adapter.md index af1e9afe97df..8647071481fd 100644 --- a/docs/source/en/api/pipelines/stable_diffusion/adapter.md +++ b/docs/source/en/api/pipelines/stable_diffusion/adapter.md @@ -29,8 +29,6 @@ This model was contributed by the community contributor [HimariO](https://github - __call__ - enable_attention_slicing - disable_attention_slicing - - enable_vae_slicing - - disable_vae_slicing - enable_xformers_memory_efficient_attention - disable_xformers_memory_efficient_attention @@ -41,7 +39,5 @@ This model was contributed by the community contributor [HimariO](https://github - __call__ - enable_attention_slicing - disable_attention_slicing - - enable_vae_slicing - - disable_vae_slicing - enable_xformers_memory_efficient_attention - disable_xformers_memory_efficient_attention diff --git a/docs/source/en/api/pipelines/stable_diffusion/text2img.md b/docs/source/en/api/pipelines/stable_diffusion/text2img.md index 2033638a1a36..f3d032320ae8 100644 --- a/docs/source/en/api/pipelines/stable_diffusion/text2img.md +++ b/docs/source/en/api/pipelines/stable_diffusion/text2img.md @@ -34,12 +34,8 @@ The abstract from the paper is: - __call__ - enable_attention_slicing - disable_attention_slicing - - enable_vae_slicing - - disable_vae_slicing - enable_xformers_memory_efficient_attention - disable_xformers_memory_efficient_attention - - enable_vae_tiling - - disable_vae_tiling - load_textual_inversion - from_single_file - load_lora_weights diff --git a/docs/source/en/api/pipelines/stable_unclip.md b/docs/source/en/api/pipelines/stable_unclip.md index 09100201bb1e..2284ee144d58 100644 --- a/docs/source/en/api/pipelines/stable_unclip.md +++ b/docs/source/en/api/pipelines/stable_unclip.md @@ -106,8 +106,6 @@ image - __call__ - enable_attention_slicing - disable_attention_slicing - - enable_vae_slicing - - disable_vae_slicing - enable_xformers_memory_efficient_attention - disable_xformers_memory_efficient_attention @@ -118,8 +116,6 @@ image - __call__ - enable_attention_slicing - disable_attention_slicing - - enable_vae_slicing - - disable_vae_slicing - enable_xformers_memory_efficient_attention - disable_xformers_memory_efficient_attention diff --git a/docs/source/en/optimization/memory.md b/docs/source/en/optimization/memory.md index 5212b70c9cea..e95fdffb9e28 100644 --- a/docs/source/en/optimization/memory.md +++ b/docs/source/en/optimization/memory.md @@ -173,7 +173,7 @@ VAE slicing saves memory by splitting large batches of inputs into a single batc For example, if you're generating 4 images at once, decoding would increase peak activation memory by 4x. VAE slicing reduces this by only decoding 1 image at a time instead of all 4 images at once. -Call [`~StableDiffusionPipeline.enable_vae_slicing`] to enable sliced VAE. You can expect a small increase in performance when decoding multi-image batches and no performance impact for single-image batches. +Call [`~AutoencoderKL.enable_slicing`] to enable sliced VAE. You can expect a small increase in performance when decoding multi-image batches and no performance impact for single-image batches. ```py import torch @@ -183,7 +183,7 @@ pipeline = StableDiffusionXLPipeline.from_pretrained( "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, ).to("cuda") -pipeline.enable_vae_slicing() +pipeline.vae.enable_slicing() pipeline(["An astronaut riding a horse on Mars"]*32).images[0] print(f"Max memory reserved: {torch.cuda.max_memory_allocated() / 1024**3:.2f} GB") ``` @@ -195,7 +195,7 @@ print(f"Max memory reserved: {torch.cuda.max_memory_allocated() / 1024**3:.2f} G VAE tiling saves memory by dividing an image into smaller overlapping tiles instead of processing the entire image at once. This also reduces peak memory usage because the GPU is only processing a tile at a time. -Call [`~StableDiffusionPipeline.enable_vae_tiling`] to enable VAE tiling. The generated image may have some tone variation from tile-to-tile because they're decoded separately, but there shouldn't be any obvious seams between the tiles. Tiling is disabled for resolutions lower than a pre-specified (but configurable) limit. For example, this limit is 512x512 for the VAE in [`StableDiffusionPipeline`]. +Call [`~AutoencoderKL.enable_tiling`] to enable VAE tiling. The generated image may have some tone variation from tile-to-tile because they're decoded separately, but there shouldn't be any obvious seams between the tiles. Tiling is disabled for resolutions lower than a pre-specified (but configurable) limit. For example, this limit is 512x512 for the VAE in [`StableDiffusionPipeline`]. ```py import torch @@ -205,7 +205,7 @@ from diffusers.utils import load_image pipeline = AutoPipelineForImage2Image.from_pretrained( "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16 ).to("cuda") -pipeline.enable_vae_tiling() +pipeline.vae.enable_tiling() init_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/img2img-sdxl-init.png") prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k" diff --git a/docs/source/en/using-diffusers/ip_adapter.md b/docs/source/en/using-diffusers/ip_adapter.md index 781eb6ce00ef..e552bd254739 100644 --- a/docs/source/en/using-diffusers/ip_adapter.md +++ b/docs/source/en/using-diffusers/ip_adapter.md @@ -180,7 +180,7 @@ scheduler = DDIMScheduler.from_pretrained( steps_offset=1, ) pipeline.scheduler = scheduler -pipeline.enable_vae_slicing() +pipeline.vae.enable_slicing() pipeline.load_ip_adapter("h94/IP-Adapter", subfolder="models", weight_name="ip-adapter_sd15.bin") pipeline.enable_model_cpu_offload() diff --git a/docs/source/ko/optimization/fp16.md b/docs/source/ko/optimization/fp16.md index 56f1330c404e..0642517c0557 100644 --- a/docs/source/ko/optimization/fp16.md +++ b/docs/source/ko/optimization/fp16.md @@ -112,7 +112,7 @@ image = pipe(prompt).images[0] 이를 [`~StableDiffusionPipeline.enable_attention_slicing`] 또는 [`~StableDiffusionPipeline.enable_xformers_memory_efficient_attention`]과 결합하여 메모리 사용을 추가로 최소화할 수 있습니다. -VAE 디코드를 한 번에 하나씩 수행하려면 추론 전에 파이프라인에서 [`~StableDiffusionPipeline.enable_vae_slicing`]을 호출합니다. 예를 들어: +VAE 디코드를 한 번에 하나씩 수행하려면 추론 전에 파이프라인에서 [`~AutoencoderKL.enable_slicing`]을 호출합니다. 예를 들어: ```Python import torch @@ -126,7 +126,7 @@ pipe = StableDiffusionPipeline.from_pretrained( pipe = pipe.to("cuda") prompt = "a photo of an astronaut riding a horse on mars" -pipe.enable_vae_slicing() +pipe.vae.enable_slicing() images = pipe([prompt] * 32).images ``` diff --git a/docs/source/ko/using-diffusers/diffedit.md b/docs/source/ko/using-diffusers/diffedit.md index edf23f0214ab..626331a05236 100644 --- a/docs/source/ko/using-diffusers/diffedit.md +++ b/docs/source/ko/using-diffusers/diffedit.md @@ -53,7 +53,7 @@ pipeline = StableDiffusionDiffEditPipeline.from_pretrained( pipeline.scheduler = DDIMScheduler.from_config(pipeline.scheduler.config) pipeline.inverse_scheduler = DDIMInverseScheduler.from_config(pipeline.scheduler.config) pipeline.enable_model_cpu_offload() -pipeline.enable_vae_slicing() +pipeline.vae.enable_slicing() ``` 수정하기 위한 이미지를 불러옵니다: @@ -169,7 +169,7 @@ pipeline = StableDiffusionDiffEditPipeline.from_pretrained( "stabilityai/stable-diffusion-2-1", torch_dtype=torch.float16, use_safetensors=True ) pipeline.enable_model_cpu_offload() -pipeline.enable_vae_slicing() +pipeline.vae.enable_slicing() @torch.no_grad() def embed_prompts(sentences, tokenizer, text_encoder, device="cuda"): diff --git a/docs/source/zh/optimization/memory.md b/docs/source/zh/optimization/memory.md index 662dcaf4bcf2..7f08f63ff1d7 100644 --- a/docs/source/zh/optimization/memory.md +++ b/docs/source/zh/optimization/memory.md @@ -171,7 +171,7 @@ VAE 切片通过将大批次输入拆分为单个数据批次并分别处理它 例如，如果您同时生成 4 个图像，解码会将峰值激活内存增加 4 倍。VAE 切片通过一次只解码 1 个图像而不是所有 4 个图像来减少这种情况。 -调用 [`~StableDiffusionPipeline.enable_vae_slicing`] 来启用切片 VAE。您可以预期在解码多图像批次时性能会有小幅提升，而在单图像批次时没有性能影响。 +调用 [`~AutoencoderKL.enable_slicing`] 来启用切片 VAE。您可以预期在解码多图像批次时性能会有小幅提升，而在单图像批次时没有性能影响。 ```py import torch @@ -181,7 +181,7 @@ pipeline = StableDiffusionXLPipeline.from_pretrained( "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, ).to("cuda") -pipeline.enable_vae_slicing() +pipeline.vae.enable_slicing() pipeline(["An astronaut riding a horse on Mars"]*32).images[0] print(f"Max memory reserved: {torch.cuda.max_memory_allocated() / 1024**3:.2f} GB") ``` @@ -193,7 +193,7 @@ print(f"Max memory reserved: {torch.cuda.max_memory_allocated() / 1024**3:.2f} G VAE 平铺通过将图像划分为较小的重叠图块而不是一次性处理整个图像来节省内存。这也减少了峰值内存使用量，因为 GPU 一次只处理一个图块。 -调用 [`~StableDiffusionPipeline.enable_vae_tiling`] 来启用 VAE 平铺。生成的图像可能因图块到图块的色调变化而有所不同，因为它们被单独解码，但图块之间不应有明显的接缝。对于低于预设（但可配置）限制的分辨率，平铺被禁用。例如，对于 [`StableDiffusionPipeline`] 中的 VAE，此限制为 512x512。 +调用 [`~AutoencoderKL.enable_tiling`] 来启用 VAE 平铺。生成的图像可能因图块到图块的色调变化而有所不同，因为它们被单独解码，但图块之间不应有明显的接缝。对于低于预设（但可配置）限制的分辨率，平铺被禁用。例如，对于 [`StableDiffusionPipeline`] 中的 VAE，此限制为 512x512。 ```py import torch @@ -203,7 +203,7 @@ from diffusers.utils import load_image pipeline = AutoPipelineForImage2Image.from_pretrained( "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16 ).to("cuda") -pipeline.enable_vae_tiling() +pipeline.vae.enable_tiling() init_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/img2img-sdxl-init.png") prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k" diff --git a/examples/advanced_diffusion_training/test_dreambooth_lora_flux_advanced.py b/examples/advanced_diffusion_training/test_dreambooth_lora_flux_advanced.py index 74aec95dd8a9..ddc863b5116c 100644 --- a/examples/advanced_diffusion_training/test_dreambooth_lora_flux_advanced.py +++ b/examples/advanced_diffusion_training/test_dreambooth_lora_flux_advanced.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/examples/advanced_diffusion_training/train_dreambooth_lora_flux_advanced.py b/examples/advanced_diffusion_training/train_dreambooth_lora_flux_advanced.py index 005f4303c3c1..53d9da3df72e 100644 --- a/examples/advanced_diffusion_training/train_dreambooth_lora_flux_advanced.py +++ b/examples/advanced_diffusion_training/train_dreambooth_lora_flux_advanced.py @@ -94,7 +94,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.39.0.dev0") +check_min_version("0.40.0.dev0") logger = get_logger(__name__) diff --git a/examples/advanced_diffusion_training/train_dreambooth_lora_sd15_advanced.py b/examples/advanced_diffusion_training/train_dreambooth_lora_sd15_advanced.py index e10e442a7d61..2eae63432a64 100644 --- a/examples/advanced_diffusion_training/train_dreambooth_lora_sd15_advanced.py +++ b/examples/advanced_diffusion_training/train_dreambooth_lora_sd15_advanced.py @@ -88,7 +88,7 @@ # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.39.0.dev0") +check_min_version("0.40.0.dev0") logger = get_logger(__name__) diff --git a/examples/advanced_diffusion_training/train_dreambooth_lora_sdxl_advanced.py b/examples/advanced_diffusion_training/train_dreambooth_lora_sdxl_advanced.py index cea4d536da95..b4b653875088 100644 --- a/examples/advanced_diffusion_training/train_dreambooth_lora_sdxl_advanced.py +++ b/examples/advanced_diffusion_training/train_dreambooth_lora_sdxl_advanced.py @@ -95,7 +95,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.39.0.dev0") +check_min_version("0.40.0.dev0") logger = get_logger(__name__) diff --git a/examples/cogvideo/train_cogvideox_image_to_video_lora.py b/examples/cogvideo/train_cogvideox_image_to_video_lora.py index 311fe0b4cf5c..8832f60e5d0a 100644 --- a/examples/cogvideo/train_cogvideox_image_to_video_lora.py +++ b/examples/cogvideo/train_cogvideox_image_to_video_lora.py @@ -61,7 +61,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.39.0.dev0") +check_min_version("0.40.0.dev0") logger = get_logger(__name__) diff --git a/examples/cogvideo/train_cogvideox_lora.py b/examples/cogvideo/train_cogvideox_lora.py index 364ed2500f03..40a55f94fa11 100644 --- a/examples/cogvideo/train_cogvideox_lora.py +++ b/examples/cogvideo/train_cogvideox_lora.py @@ -52,7 +52,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.39.0.dev0") +check_min_version("0.40.0.dev0") logger = get_logger(__name__) diff --git a/examples/cogview4-control/requirements.txt b/examples/cogview4-control/requirements.txt index 6c5ec2e03f9a..7eabb2592c9a 100644 --- a/examples/cogview4-control/requirements.txt +++ b/examples/cogview4-control/requirements.txt @@ -1,4 +1,4 @@ -transformers==4.47.0 +transformers==5.3.0 wandb torch torchvision diff --git a/examples/cogview4-control/train_control_cogview4.py b/examples/cogview4-control/train_control_cogview4.py index 7aee41e460c3..30d8f234d79b 100644 --- a/examples/cogview4-control/train_control_cogview4.py +++ b/examples/cogview4-control/train_control_cogview4.py @@ -60,7 +60,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.39.0.dev0") +check_min_version("0.40.0.dev0") logger = get_logger(__name__) diff --git a/examples/community/marigold_depth_estimation.py b/examples/community/marigold_depth_estimation.py index f619cef19a17..3d5bc03ee088 100644 --- a/examples/community/marigold_depth_estimation.py +++ b/examples/community/marigold_depth_estimation.py @@ -43,7 +43,7 @@ # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.39.0.dev0") +check_min_version("0.40.0.dev0") class MarigoldDepthOutput(BaseOutput): diff --git a/examples/consistency_distillation/test_lcm_lora.py b/examples/consistency_distillation/test_lcm_lora.py index 1eeb31d6e414..37ce510db3ee 100644 --- a/examples/consistency_distillation/test_lcm_lora.py +++ b/examples/consistency_distillation/test_lcm_lora.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/examples/consistency_distillation/train_lcm_distill_lora_sd_wds.py b/examples/consistency_distillation/train_lcm_distill_lora_sd_wds.py index dc7c0b5bcbb6..f1479d50b408 100644 --- a/examples/consistency_distillation/train_lcm_distill_lora_sd_wds.py +++ b/examples/consistency_distillation/train_lcm_distill_lora_sd_wds.py @@ -74,7 +74,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.39.0.dev0") +check_min_version("0.40.0.dev0") logger = get_logger(__name__) diff --git a/examples/consistency_distillation/train_lcm_distill_lora_sdxl.py b/examples/consistency_distillation/train_lcm_distill_lora_sdxl.py index a350910fb226..cd9f6c1b7aa1 100644 --- a/examples/consistency_distillation/train_lcm_distill_lora_sdxl.py +++ b/examples/consistency_distillation/train_lcm_distill_lora_sdxl.py @@ -67,7 +67,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.39.0.dev0") +check_min_version("0.40.0.dev0") logger = get_logger(__name__) diff --git a/examples/consistency_distillation/train_lcm_distill_lora_sdxl_wds.py b/examples/consistency_distillation/train_lcm_distill_lora_sdxl_wds.py index 82a6330f6686..99c67c042059 100644 --- a/examples/consistency_distillation/train_lcm_distill_lora_sdxl_wds.py +++ b/examples/consistency_distillation/train_lcm_distill_lora_sdxl_wds.py @@ -80,7 +80,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.39.0.dev0") +check_min_version("0.40.0.dev0") logger = get_logger(__name__) diff --git a/examples/consistency_distillation/train_lcm_distill_sd_wds.py b/examples/consistency_distillation/train_lcm_distill_sd_wds.py index a5e4df573d1e..988c12aa101b 100644 --- a/examples/consistency_distillation/train_lcm_distill_sd_wds.py +++ b/examples/consistency_distillation/train_lcm_distill_sd_wds.py @@ -73,7 +73,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.39.0.dev0") +check_min_version("0.40.0.dev0") logger = get_logger(__name__) diff --git a/examples/consistency_distillation/train_lcm_distill_sdxl_wds.py b/examples/consistency_distillation/train_lcm_distill_sdxl_wds.py index 4149158ded90..53030a98c37d 100644 --- a/examples/consistency_distillation/train_lcm_distill_sdxl_wds.py +++ b/examples/consistency_distillation/train_lcm_distill_sdxl_wds.py @@ -79,7 +79,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.39.0.dev0") +check_min_version("0.40.0.dev0") logger = get_logger(__name__) diff --git a/examples/controlnet/test_controlnet.py b/examples/controlnet/test_controlnet.py index 04b9a950a97c..6085e20b23bd 100644 --- a/examples/controlnet/test_controlnet.py +++ b/examples/controlnet/test_controlnet.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/examples/controlnet/train_controlnet.py b/examples/controlnet/train_controlnet.py index 515d6b0d18d5..ff6bdc1e05b4 100644 --- a/examples/controlnet/train_controlnet.py +++ b/examples/controlnet/train_controlnet.py @@ -61,7 +61,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.39.0.dev0") +check_min_version("0.40.0.dev0") logger = get_logger(__name__) diff --git a/examples/controlnet/train_controlnet_flax.py b/examples/controlnet/train_controlnet_flax.py index 76bb2959123e..5322f8d64d6a 100644 --- a/examples/controlnet/train_controlnet_flax.py +++ b/examples/controlnet/train_controlnet_flax.py @@ -61,7 +61,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.39.0.dev0") +check_min_version("0.40.0.dev0") logger = logging.getLogger(__name__) diff --git a/examples/controlnet/train_controlnet_flux.py b/examples/controlnet/train_controlnet_flux.py index c06f98acb89c..5b44c0b1db79 100644 --- a/examples/controlnet/train_controlnet_flux.py +++ b/examples/controlnet/train_controlnet_flux.py @@ -66,7 +66,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.39.0.dev0") +check_min_version("0.40.0.dev0") logger = get_logger(__name__) if is_torch_npu_available(): diff --git a/examples/controlnet/train_controlnet_sd3.py b/examples/controlnet/train_controlnet_sd3.py index 19fba1cd6b0d..b8ed6c40e0f0 100644 --- a/examples/controlnet/train_controlnet_sd3.py +++ b/examples/controlnet/train_controlnet_sd3.py @@ -63,7 +63,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.39.0.dev0") +check_min_version("0.40.0.dev0") logger = get_logger(__name__) diff --git a/examples/controlnet/train_controlnet_sdxl.py b/examples/controlnet/train_controlnet_sdxl.py index 3404a857e773..2c6f71a0e71e 100644 --- a/examples/controlnet/train_controlnet_sdxl.py +++ b/examples/controlnet/train_controlnet_sdxl.py @@ -62,7 +62,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.39.0.dev0") +check_min_version("0.40.0.dev0") logger = get_logger(__name__) if is_torch_npu_available(): diff --git a/examples/custom_diffusion/test_custom_diffusion.py b/examples/custom_diffusion/test_custom_diffusion.py index ad18eb246777..dbaca56302e9 100644 --- a/examples/custom_diffusion/test_custom_diffusion.py +++ b/examples/custom_diffusion/test_custom_diffusion.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/examples/custom_diffusion/train_custom_diffusion.py b/examples/custom_diffusion/train_custom_diffusion.py index 4c6b63744657..8fe86ee42adf 100644 --- a/examples/custom_diffusion/train_custom_diffusion.py +++ b/examples/custom_diffusion/train_custom_diffusion.py @@ -64,7 +64,7 @@ # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.39.0.dev0") +check_min_version("0.40.0.dev0") logger = get_logger(__name__) diff --git a/examples/dreambooth/test_dreambooth.py b/examples/dreambooth/test_dreambooth.py index 76f7fe8bd4a5..0cc637472841 100644 --- a/examples/dreambooth/test_dreambooth.py +++ b/examples/dreambooth/test_dreambooth.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/examples/dreambooth/test_dreambooth_flux.py b/examples/dreambooth/test_dreambooth_flux.py index 4b9117266cca..93b7883608ea 100644 --- a/examples/dreambooth/test_dreambooth_flux.py +++ b/examples/dreambooth/test_dreambooth_flux.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/examples/dreambooth/test_dreambooth_lora.py b/examples/dreambooth/test_dreambooth_lora.py index e950807d372d..2c9ff897d692 100644 --- a/examples/dreambooth/test_dreambooth_lora.py +++ b/examples/dreambooth/test_dreambooth_lora.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/examples/dreambooth/test_dreambooth_lora_edm.py b/examples/dreambooth/test_dreambooth_lora_edm.py index 737f5b4309e1..c2405a2b20e8 100644 --- a/examples/dreambooth/test_dreambooth_lora_edm.py +++ b/examples/dreambooth/test_dreambooth_lora_edm.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/examples/dreambooth/test_dreambooth_lora_flux.py b/examples/dreambooth/test_dreambooth_lora_flux.py index 3b1391cb3523..6235fdc09ea3 100644 --- a/examples/dreambooth/test_dreambooth_lora_flux.py +++ b/examples/dreambooth/test_dreambooth_lora_flux.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/examples/dreambooth/test_dreambooth_lora_flux2.py b/examples/dreambooth/test_dreambooth_lora_flux2.py index 80a0b502f9a2..bc493b5610b5 100644 --- a/examples/dreambooth/test_dreambooth_lora_flux2.py +++ b/examples/dreambooth/test_dreambooth_lora_flux2.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/examples/dreambooth/test_dreambooth_lora_flux2_klein.py b/examples/dreambooth/test_dreambooth_lora_flux2_klein.py index 0e5506e1a3eb..1bc8da075374 100644 --- a/examples/dreambooth/test_dreambooth_lora_flux2_klein.py +++ b/examples/dreambooth/test_dreambooth_lora_flux2_klein.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/examples/dreambooth/test_dreambooth_lora_flux_kontext.py b/examples/dreambooth/test_dreambooth_lora_flux_kontext.py index c12fdd79ee57..4795f385a4a2 100644 --- a/examples/dreambooth/test_dreambooth_lora_flux_kontext.py +++ b/examples/dreambooth/test_dreambooth_lora_flux_kontext.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/examples/dreambooth/test_dreambooth_lora_hidream.py b/examples/dreambooth/test_dreambooth_lora_hidream.py index df4c70e2e86f..762217e1ddf0 100644 --- a/examples/dreambooth/test_dreambooth_lora_hidream.py +++ b/examples/dreambooth/test_dreambooth_lora_hidream.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/examples/dreambooth/test_dreambooth_lora_lumina2.py b/examples/dreambooth/test_dreambooth_lora_lumina2.py index 984ee0331ed6..9f2c81e6e877 100644 --- a/examples/dreambooth/test_dreambooth_lora_lumina2.py +++ b/examples/dreambooth/test_dreambooth_lora_lumina2.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/examples/dreambooth/test_dreambooth_lora_qwenimage.py b/examples/dreambooth/test_dreambooth_lora_qwenimage.py index 418ffd1bc027..0a040d999a6f 100644 --- a/examples/dreambooth/test_dreambooth_lora_qwenimage.py +++ b/examples/dreambooth/test_dreambooth_lora_qwenimage.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/examples/dreambooth/test_dreambooth_lora_sana.py b/examples/dreambooth/test_dreambooth_lora_sana.py index 7564ab08b028..57939096c184 100644 --- a/examples/dreambooth/test_dreambooth_lora_sana.py +++ b/examples/dreambooth/test_dreambooth_lora_sana.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/examples/dreambooth/test_dreambooth_lora_sd3.py b/examples/dreambooth/test_dreambooth_lora_sd3.py index 134aeeb1da66..47fedb21c70e 100644 --- a/examples/dreambooth/test_dreambooth_lora_sd3.py +++ b/examples/dreambooth/test_dreambooth_lora_sd3.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/examples/dreambooth/test_dreambooth_sd3.py b/examples/dreambooth/test_dreambooth_sd3.py index d1b3daf94427..d8f6f09cb389 100644 --- a/examples/dreambooth/test_dreambooth_sd3.py +++ b/examples/dreambooth/test_dreambooth_sd3.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/examples/dreambooth/train_dreambooth.py b/examples/dreambooth/train_dreambooth.py index 7d9af890d25f..c1f700f9d150 100644 --- a/examples/dreambooth/train_dreambooth.py +++ b/examples/dreambooth/train_dreambooth.py @@ -64,7 +64,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.39.0.dev0") +check_min_version("0.40.0.dev0") logger = get_logger(__name__) diff --git a/examples/dreambooth/train_dreambooth_flax.py b/examples/dreambooth/train_dreambooth_flax.py index b281a02f20e2..0204ee82c368 100644 --- a/examples/dreambooth/train_dreambooth_flax.py +++ b/examples/dreambooth/train_dreambooth_flax.py @@ -35,7 +35,7 @@ # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.39.0.dev0") +check_min_version("0.40.0.dev0") # Cache compiled models across invocations of this script. cc.initialize_cache(os.path.expanduser("~/.cache/jax/compilation_cache")) diff --git a/examples/dreambooth/train_dreambooth_flux.py b/examples/dreambooth/train_dreambooth_flux.py index 89e1c9dc57ad..a6b8e553249a 100644 --- a/examples/dreambooth/train_dreambooth_flux.py +++ b/examples/dreambooth/train_dreambooth_flux.py @@ -80,7 +80,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.39.0.dev0") +check_min_version("0.40.0.dev0") logger = get_logger(__name__) diff --git a/examples/dreambooth/train_dreambooth_lora.py b/examples/dreambooth/train_dreambooth_lora.py index 0ee636b4a00e..e67e64afc832 100644 --- a/examples/dreambooth/train_dreambooth_lora.py +++ b/examples/dreambooth/train_dreambooth_lora.py @@ -75,7 +75,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.39.0.dev0") +check_min_version("0.40.0.dev0") logger = get_logger(__name__) diff --git a/examples/dreambooth/train_dreambooth_lora_flux.py b/examples/dreambooth/train_dreambooth_lora_flux.py index 5fb666a4d42c..f70f53704c86 100644 --- a/examples/dreambooth/train_dreambooth_lora_flux.py +++ b/examples/dreambooth/train_dreambooth_lora_flux.py @@ -92,7 +92,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.39.0.dev0") +check_min_version("0.40.0.dev0") logger = get_logger(__name__) diff --git a/examples/dreambooth/train_dreambooth_lora_flux2.py b/examples/dreambooth/train_dreambooth_lora_flux2.py index 886e251937e6..5643fb95a910 100644 --- a/examples/dreambooth/train_dreambooth_lora_flux2.py +++ b/examples/dreambooth/train_dreambooth_lora_flux2.py @@ -104,7 +104,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.39.0.dev0") +check_min_version("0.40.0.dev0") logger = get_logger(__name__) diff --git a/examples/dreambooth/train_dreambooth_lora_flux2_img2img.py b/examples/dreambooth/train_dreambooth_lora_flux2_img2img.py index 477697fadb64..514832dd31db 100644 --- a/examples/dreambooth/train_dreambooth_lora_flux2_img2img.py +++ b/examples/dreambooth/train_dreambooth_lora_flux2_img2img.py @@ -104,7 +104,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.39.0.dev0") +check_min_version("0.40.0.dev0") logger = get_logger(__name__) diff --git a/examples/dreambooth/train_dreambooth_lora_flux2_klein.py b/examples/dreambooth/train_dreambooth_lora_flux2_klein.py index 7eb627e4bd1d..7b7e1634f8a7 100644 --- a/examples/dreambooth/train_dreambooth_lora_flux2_klein.py +++ b/examples/dreambooth/train_dreambooth_lora_flux2_klein.py @@ -104,7 +104,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.39.0.dev0") +check_min_version("0.40.0.dev0") logger = get_logger(__name__) diff --git a/examples/dreambooth/train_dreambooth_lora_flux2_klein_img2img.py b/examples/dreambooth/train_dreambooth_lora_flux2_klein_img2img.py index cdc1de9c3806..ca29cf2dbaff 100644 --- a/examples/dreambooth/train_dreambooth_lora_flux2_klein_img2img.py +++ b/examples/dreambooth/train_dreambooth_lora_flux2_klein_img2img.py @@ -104,7 +104,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.39.0.dev0") +check_min_version("0.40.0.dev0") logger = get_logger(__name__) diff --git a/examples/dreambooth/train_dreambooth_lora_flux_kontext.py b/examples/dreambooth/train_dreambooth_lora_flux_kontext.py index 97e0414635fb..846ce6336e54 100644 --- a/examples/dreambooth/train_dreambooth_lora_flux_kontext.py +++ b/examples/dreambooth/train_dreambooth_lora_flux_kontext.py @@ -92,7 +92,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.39.0.dev0") +check_min_version("0.40.0.dev0") logger = get_logger(__name__) diff --git a/examples/dreambooth/train_dreambooth_lora_hidream.py b/examples/dreambooth/train_dreambooth_lora_hidream.py index c87d96366c6d..9a027ad51eaa 100644 --- a/examples/dreambooth/train_dreambooth_lora_hidream.py +++ b/examples/dreambooth/train_dreambooth_lora_hidream.py @@ -75,7 +75,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.39.0.dev0") +check_min_version("0.40.0.dev0") logger = get_logger(__name__) diff --git a/examples/dreambooth/train_dreambooth_lora_ideogram4.py b/examples/dreambooth/train_dreambooth_lora_ideogram4.py index 64c63b1f2ea0..b59d04dfc643 100644 --- a/examples/dreambooth/train_dreambooth_lora_ideogram4.py +++ b/examples/dreambooth/train_dreambooth_lora_ideogram4.py @@ -102,7 +102,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.39.0.dev0") +check_min_version("0.40.0.dev0") logger = get_logger(__name__) diff --git a/examples/dreambooth/train_dreambooth_lora_krea2.py b/examples/dreambooth/train_dreambooth_lora_krea2.py index 01c303dd0f25..0ac631ea0f54 100644 --- a/examples/dreambooth/train_dreambooth_lora_krea2.py +++ b/examples/dreambooth/train_dreambooth_lora_krea2.py @@ -93,7 +93,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.39.0.dev0") +check_min_version("0.40.0.dev0") logger = get_logger(__name__) diff --git a/examples/dreambooth/train_dreambooth_lora_lumina2.py b/examples/dreambooth/train_dreambooth_lora_lumina2.py index 2f744fd9cc6b..381997e9a294 100644 --- a/examples/dreambooth/train_dreambooth_lora_lumina2.py +++ b/examples/dreambooth/train_dreambooth_lora_lumina2.py @@ -73,7 +73,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.39.0.dev0") +check_min_version("0.40.0.dev0") logger = get_logger(__name__) diff --git a/examples/dreambooth/train_dreambooth_lora_qwen_image.py b/examples/dreambooth/train_dreambooth_lora_qwen_image.py index 573e0bf53f8a..2befd22f3192 100644 --- a/examples/dreambooth/train_dreambooth_lora_qwen_image.py +++ b/examples/dreambooth/train_dreambooth_lora_qwen_image.py @@ -93,7 +93,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.39.0.dev0") +check_min_version("0.40.0.dev0") logger = get_logger(__name__) diff --git a/examples/dreambooth/train_dreambooth_lora_sana.py b/examples/dreambooth/train_dreambooth_lora_sana.py index 29d284611a0d..7527290e9dce 100644 --- a/examples/dreambooth/train_dreambooth_lora_sana.py +++ b/examples/dreambooth/train_dreambooth_lora_sana.py @@ -91,7 +91,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.39.0.dev0") +check_min_version("0.40.0.dev0") logger = get_logger(__name__) diff --git a/examples/dreambooth/train_dreambooth_lora_sd3.py b/examples/dreambooth/train_dreambooth_lora_sd3.py index 396f18113bf5..2c94d5e6e0e8 100644 --- a/examples/dreambooth/train_dreambooth_lora_sd3.py +++ b/examples/dreambooth/train_dreambooth_lora_sd3.py @@ -73,7 +73,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.39.0.dev0") +check_min_version("0.40.0.dev0") logger = get_logger(__name__) diff --git a/examples/dreambooth/train_dreambooth_lora_sdxl.py b/examples/dreambooth/train_dreambooth_lora_sdxl.py index ac8dd9243df6..43127665cd06 100644 --- a/examples/dreambooth/train_dreambooth_lora_sdxl.py +++ b/examples/dreambooth/train_dreambooth_lora_sdxl.py @@ -80,7 +80,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.39.0.dev0") +check_min_version("0.40.0.dev0") logger = get_logger(__name__) diff --git a/examples/dreambooth/train_dreambooth_lora_z_image.py b/examples/dreambooth/train_dreambooth_lora_z_image.py index a54c84b0798f..de3e01e4f2de 100644 --- a/examples/dreambooth/train_dreambooth_lora_z_image.py +++ b/examples/dreambooth/train_dreambooth_lora_z_image.py @@ -104,7 +104,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.39.0.dev0") +check_min_version("0.40.0.dev0") logger = get_logger(__name__) diff --git a/examples/dreambooth/train_dreambooth_sd3.py b/examples/dreambooth/train_dreambooth_sd3.py index d7dfebe7133f..2897d0179a6b 100644 --- a/examples/dreambooth/train_dreambooth_sd3.py +++ b/examples/dreambooth/train_dreambooth_sd3.py @@ -64,7 +64,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.39.0.dev0") +check_min_version("0.40.0.dev0") logger = get_logger(__name__) diff --git a/examples/flux-control/train_control_flux.py b/examples/flux-control/train_control_flux.py index fb5edd185b6f..9dcff777ddd6 100644 --- a/examples/flux-control/train_control_flux.py +++ b/examples/flux-control/train_control_flux.py @@ -55,7 +55,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.39.0.dev0") +check_min_version("0.40.0.dev0") logger = get_logger(__name__) diff --git a/examples/flux-control/train_control_lora_flux.py b/examples/flux-control/train_control_lora_flux.py index 3e0c2ee64393..b12930e3e87f 100644 --- a/examples/flux-control/train_control_lora_flux.py +++ b/examples/flux-control/train_control_lora_flux.py @@ -58,7 +58,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.39.0.dev0") +check_min_version("0.40.0.dev0") logger = get_logger(__name__) diff --git a/examples/instruct_pix2pix/test_instruct_pix2pix.py b/examples/instruct_pix2pix/test_instruct_pix2pix.py index cfb9c8d8540e..14a7137d6050 100644 --- a/examples/instruct_pix2pix/test_instruct_pix2pix.py +++ b/examples/instruct_pix2pix/test_instruct_pix2pix.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/examples/instruct_pix2pix/train_instruct_pix2pix.py b/examples/instruct_pix2pix/train_instruct_pix2pix.py index 89eb2504e97a..4b860823990e 100644 --- a/examples/instruct_pix2pix/train_instruct_pix2pix.py +++ b/examples/instruct_pix2pix/train_instruct_pix2pix.py @@ -58,7 +58,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.39.0.dev0") +check_min_version("0.40.0.dev0") logger = get_logger(__name__, log_level="INFO") diff --git a/examples/instruct_pix2pix/train_instruct_pix2pix_sdxl.py b/examples/instruct_pix2pix/train_instruct_pix2pix_sdxl.py index 4b74e3b61607..d041f53964eb 100644 --- a/examples/instruct_pix2pix/train_instruct_pix2pix_sdxl.py +++ b/examples/instruct_pix2pix/train_instruct_pix2pix_sdxl.py @@ -60,7 +60,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.39.0.dev0") +check_min_version("0.40.0.dev0") logger = get_logger(__name__, log_level="INFO") diff --git a/examples/kandinsky2_2/text_to_image/train_text_to_image_decoder.py b/examples/kandinsky2_2/text_to_image/train_text_to_image_decoder.py index 73b3856ccb3f..69934c0625f3 100644 --- a/examples/kandinsky2_2/text_to_image/train_text_to_image_decoder.py +++ b/examples/kandinsky2_2/text_to_image/train_text_to_image_decoder.py @@ -53,7 +53,7 @@ # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.39.0.dev0") +check_min_version("0.40.0.dev0") logger = get_logger(__name__, log_level="INFO") diff --git a/examples/kandinsky2_2/text_to_image/train_text_to_image_lora_decoder.py b/examples/kandinsky2_2/text_to_image/train_text_to_image_lora_decoder.py index 3e7eb84d9318..1b447c3e2fef 100644 --- a/examples/kandinsky2_2/text_to_image/train_text_to_image_lora_decoder.py +++ b/examples/kandinsky2_2/text_to_image/train_text_to_image_lora_decoder.py @@ -46,7 +46,7 @@ # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.39.0.dev0") +check_min_version("0.40.0.dev0") logger = get_logger(__name__, log_level="INFO") diff --git a/examples/kandinsky2_2/text_to_image/train_text_to_image_lora_prior.py b/examples/kandinsky2_2/text_to_image/train_text_to_image_lora_prior.py index 185bd0709875..dbfd4f370c44 100644 --- a/examples/kandinsky2_2/text_to_image/train_text_to_image_lora_prior.py +++ b/examples/kandinsky2_2/text_to_image/train_text_to_image_lora_prior.py @@ -46,7 +46,7 @@ # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.39.0.dev0") +check_min_version("0.40.0.dev0") logger = get_logger(__name__, log_level="INFO") diff --git a/examples/kandinsky2_2/text_to_image/train_text_to_image_prior.py b/examples/kandinsky2_2/text_to_image/train_text_to_image_prior.py index 51a847e1d842..4ab539b54949 100644 --- a/examples/kandinsky2_2/text_to_image/train_text_to_image_prior.py +++ b/examples/kandinsky2_2/text_to_image/train_text_to_image_prior.py @@ -52,7 +52,7 @@ # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.39.0.dev0") +check_min_version("0.40.0.dev0") logger = get_logger(__name__, log_level="INFO") diff --git a/examples/t2i_adapter/test_t2i_adapter.py b/examples/t2i_adapter/test_t2i_adapter.py index be7331d024ca..d347434018d0 100644 --- a/examples/t2i_adapter/test_t2i_adapter.py +++ b/examples/t2i_adapter/test_t2i_adapter.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/examples/t2i_adapter/train_t2i_adapter_sdxl.py b/examples/t2i_adapter/train_t2i_adapter_sdxl.py index 0e47546cf68a..32863d469a6c 100644 --- a/examples/t2i_adapter/train_t2i_adapter_sdxl.py +++ b/examples/t2i_adapter/train_t2i_adapter_sdxl.py @@ -61,7 +61,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.39.0.dev0") +check_min_version("0.40.0.dev0") logger = get_logger(__name__) diff --git a/examples/test_examples_utils.py b/examples/test_examples_utils.py index f3f3d7541cb2..0a4d9687ccf5 100644 --- a/examples/test_examples_utils.py +++ b/examples/test_examples_utils.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/examples/text_to_image/train_text_to_image.py b/examples/text_to_image/train_text_to_image.py index 0c15090f3a49..77a4c6237c55 100644 --- a/examples/text_to_image/train_text_to_image.py +++ b/examples/text_to_image/train_text_to_image.py @@ -57,7 +57,7 @@ # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.39.0.dev0") +check_min_version("0.40.0.dev0") logger = get_logger(__name__, log_level="INFO") diff --git a/examples/text_to_image/train_text_to_image_flax.py b/examples/text_to_image/train_text_to_image_flax.py index 8f973d2e4401..7a93713eaf04 100644 --- a/examples/text_to_image/train_text_to_image_flax.py +++ b/examples/text_to_image/train_text_to_image_flax.py @@ -49,7 +49,7 @@ # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.39.0.dev0") +check_min_version("0.40.0.dev0") logger = logging.getLogger(__name__) diff --git a/examples/text_to_image/train_text_to_image_lora.py b/examples/text_to_image/train_text_to_image_lora.py index bd9064202308..d33908b28000 100644 --- a/examples/text_to_image/train_text_to_image_lora.py +++ b/examples/text_to_image/train_text_to_image_lora.py @@ -61,7 +61,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.39.0.dev0") +check_min_version("0.40.0.dev0") logger = get_logger(__name__, log_level="INFO") diff --git a/examples/text_to_image/train_text_to_image_lora_sdxl.py b/examples/text_to_image/train_text_to_image_lora_sdxl.py index 0996cf8cc5cd..eafc03d69867 100644 --- a/examples/text_to_image/train_text_to_image_lora_sdxl.py +++ b/examples/text_to_image/train_text_to_image_lora_sdxl.py @@ -68,7 +68,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.39.0.dev0") +check_min_version("0.40.0.dev0") logger = get_logger(__name__) if is_torch_npu_available(): diff --git a/examples/text_to_image/train_text_to_image_sdxl.py b/examples/text_to_image/train_text_to_image_sdxl.py index 8eef6410cf5d..70bf0714573e 100644 --- a/examples/text_to_image/train_text_to_image_sdxl.py +++ b/examples/text_to_image/train_text_to_image_sdxl.py @@ -55,7 +55,7 @@ # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.39.0.dev0") +check_min_version("0.40.0.dev0") logger = get_logger(__name__) if is_torch_npu_available(): diff --git a/examples/textual_inversion/test_textual_inversion.py b/examples/textual_inversion/test_textual_inversion.py index baf8692f2275..27dddc3092a9 100644 --- a/examples/textual_inversion/test_textual_inversion.py +++ b/examples/textual_inversion/test_textual_inversion.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/examples/textual_inversion/test_textual_inversion_sdxl.py b/examples/textual_inversion/test_textual_inversion_sdxl.py index 3af75b44ee5f..1d0f8f26c477 100644 --- a/examples/textual_inversion/test_textual_inversion_sdxl.py +++ b/examples/textual_inversion/test_textual_inversion_sdxl.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/examples/textual_inversion/textual_inversion.py b/examples/textual_inversion/textual_inversion.py index 24a3bda2f49b..a6f04399e872 100644 --- a/examples/textual_inversion/textual_inversion.py +++ b/examples/textual_inversion/textual_inversion.py @@ -82,7 +82,7 @@ # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.39.0.dev0") +check_min_version("0.40.0.dev0") logger = get_logger(__name__) diff --git a/examples/textual_inversion/textual_inversion_flax.py b/examples/textual_inversion/textual_inversion_flax.py index 54cebf646da7..8ff75c60e176 100644 --- a/examples/textual_inversion/textual_inversion_flax.py +++ b/examples/textual_inversion/textual_inversion_flax.py @@ -56,7 +56,7 @@ # ------------------------------------------------------------------------------ # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.39.0.dev0") +check_min_version("0.40.0.dev0") logger = logging.getLogger(__name__) diff --git a/examples/textual_inversion/textual_inversion_sdxl.py b/examples/textual_inversion/textual_inversion_sdxl.py index 3a77c3e3b071..99834ffae103 100644 --- a/examples/textual_inversion/textual_inversion_sdxl.py +++ b/examples/textual_inversion/textual_inversion_sdxl.py @@ -77,7 +77,7 @@ # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.39.0.dev0") +check_min_version("0.40.0.dev0") logger = get_logger(__name__) diff --git a/examples/unconditional_image_generation/test_unconditional.py b/examples/unconditional_image_generation/test_unconditional.py index 94ea88881e52..cb984e83f82c 100644 --- a/examples/unconditional_image_generation/test_unconditional.py +++ b/examples/unconditional_image_generation/test_unconditional.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/examples/unconditional_image_generation/train_unconditional.py b/examples/unconditional_image_generation/train_unconditional.py index bd981688bae2..cf8dd2cce9ae 100644 --- a/examples/unconditional_image_generation/train_unconditional.py +++ b/examples/unconditional_image_generation/train_unconditional.py @@ -29,7 +29,7 @@ # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.39.0.dev0") +check_min_version("0.40.0.dev0") logger = get_logger(__name__, log_level="INFO") diff --git a/examples/vqgan/train_vqgan.py b/examples/vqgan/train_vqgan.py index b3e6b1889153..95af9cbf8b21 100644 --- a/examples/vqgan/train_vqgan.py +++ b/examples/vqgan/train_vqgan.py @@ -50,7 +50,7 @@ import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. -check_min_version("0.39.0.dev0") +check_min_version("0.40.0.dev0") logger = get_logger(__name__, log_level="INFO") diff --git a/scripts/convert_joyimage_edit_to_diffusers.py b/scripts/convert_joyimage_edit_to_diffusers.py index 3ad23de8f462..3fe4a7b12cd3 100644 --- a/scripts/convert_joyimage_edit_to_diffusers.py +++ b/scripts/convert_joyimage_edit_to_diffusers.py @@ -1,5 +1,29 @@ +"""Convert JoyImage Edit / Edit Plus checkpoints to diffusers format. + +Supports both JoyImage-Edit (single-image editing) and JoyImage-Edit-Plus +(multi-image editing). The transformer weight layout is identical; only the +target model class and pipeline differ. + +Usage: + # Convert JoyImage Edit (default) + python convert_joyimage_edit_to_diffusers.py \ + --transformer_ckpt_path /path/to/transformer.pt \ + --vae_ckpt_path /path/to/vae.pt \ + --text_encoder_path Qwen/Qwen3-VL-8B-Instruct \ + --output_path /path/to/output \ + --save_pipeline + + # Convert JoyImage Edit Plus + python convert_joyimage_edit_to_diffusers.py \ + --model_type edit_plus \ + --transformer_ckpt_path /path/to/transformer.pt \ + --vae_ckpt_path /path/to/vae.pt \ + --text_encoder_path Qwen/Qwen3-VL-8B-Instruct \ + --output_path /path/to/output \ + --save_pipeline +""" + import argparse -from typing import Any, Dict, Tuple import torch from accelerate import init_empty_weights @@ -10,19 +34,18 @@ JoyImageEditPipeline, JoyImageEditTransformer3DModel, ) +from diffusers.models.transformers.transformer_joyimage_edit_plus import JoyImageEditPlusTransformer3DModel +from diffusers.pipelines.joyimage.pipeline_joyimage_edit_plus import JoyImageEditPlusPipeline from diffusers.schedulers.scheduling_flow_match_euler_discrete import ( FlowMatchEulerDiscreteScheduler, ) -# This code is modified from convert_wan_to_diffusers.py to support input ckpt path def convert_vae(vae_ckpt_path): old_state_dict = torch.load(vae_ckpt_path, weights_only=True) new_state_dict = {} - # Create mappings for specific components middle_key_mapping = { - # Encoder middle block "encoder.middle.0.residual.0.gamma": "encoder.mid_block.resnets.0.norm1.gamma", "encoder.middle.0.residual.2.bias": "encoder.mid_block.resnets.0.conv1.bias", "encoder.middle.0.residual.2.weight": "encoder.mid_block.resnets.0.conv1.weight", @@ -35,7 +58,6 @@ def convert_vae(vae_ckpt_path): "encoder.middle.2.residual.3.gamma": "encoder.mid_block.resnets.1.norm2.gamma", "encoder.middle.2.residual.6.bias": "encoder.mid_block.resnets.1.conv2.bias", "encoder.middle.2.residual.6.weight": "encoder.mid_block.resnets.1.conv2.weight", - # Decoder middle block "decoder.middle.0.residual.0.gamma": "decoder.mid_block.resnets.0.norm1.gamma", "decoder.middle.0.residual.2.bias": "decoder.mid_block.resnets.0.conv1.bias", "decoder.middle.0.residual.2.weight": "decoder.mid_block.resnets.0.conv1.weight", @@ -50,15 +72,12 @@ def convert_vae(vae_ckpt_path): "decoder.middle.2.residual.6.weight": "decoder.mid_block.resnets.1.conv2.weight", } - # Create a mapping for attention blocks attention_mapping = { - # Encoder middle attention "encoder.middle.1.norm.gamma": "encoder.mid_block.attentions.0.norm.gamma", "encoder.middle.1.to_qkv.weight": "encoder.mid_block.attentions.0.to_qkv.weight", "encoder.middle.1.to_qkv.bias": "encoder.mid_block.attentions.0.to_qkv.bias", "encoder.middle.1.proj.weight": "encoder.mid_block.attentions.0.proj.weight", "encoder.middle.1.proj.bias": "encoder.mid_block.attentions.0.proj.bias", - # Decoder middle attention "decoder.middle.1.norm.gamma": "decoder.mid_block.attentions.0.norm.gamma", "decoder.middle.1.to_qkv.weight": "decoder.mid_block.attentions.0.to_qkv.weight", "decoder.middle.1.to_qkv.bias": "decoder.mid_block.attentions.0.to_qkv.bias", @@ -66,19 +85,15 @@ def convert_vae(vae_ckpt_path): "decoder.middle.1.proj.bias": "decoder.mid_block.attentions.0.proj.bias", } - # Create a mapping for the head components head_mapping = { - # Encoder head "encoder.head.0.gamma": "encoder.norm_out.gamma", "encoder.head.2.bias": "encoder.conv_out.bias", "encoder.head.2.weight": "encoder.conv_out.weight", - # Decoder head "decoder.head.0.gamma": "decoder.norm_out.gamma", "decoder.head.2.bias": "decoder.conv_out.bias", "decoder.head.2.weight": "decoder.conv_out.weight", } - # Create a mapping for the quant components quant_mapping = { "conv1.weight": "quant_conv.weight", "conv1.bias": "quant_conv.bias", @@ -86,40 +101,25 @@ def convert_vae(vae_ckpt_path): "conv2.bias": "post_quant_conv.bias", } - # Process each key in the state dict for key, value in old_state_dict.items(): - # Handle middle block keys using the mapping if key in middle_key_mapping: - new_key = middle_key_mapping[key] - new_state_dict[new_key] = value - # Handle attention blocks using the mapping + new_state_dict[middle_key_mapping[key]] = value elif key in attention_mapping: - new_key = attention_mapping[key] - new_state_dict[new_key] = value - # Handle head keys using the mapping + new_state_dict[attention_mapping[key]] = value elif key in head_mapping: - new_key = head_mapping[key] - new_state_dict[new_key] = value - # Handle quant keys using the mapping + new_state_dict[head_mapping[key]] = value elif key in quant_mapping: - new_key = quant_mapping[key] - new_state_dict[new_key] = value - # Handle encoder conv1 + new_state_dict[quant_mapping[key]] = value elif key == "encoder.conv1.weight": new_state_dict["encoder.conv_in.weight"] = value elif key == "encoder.conv1.bias": new_state_dict["encoder.conv_in.bias"] = value - # Handle decoder conv1 elif key == "decoder.conv1.weight": new_state_dict["decoder.conv_in.weight"] = value elif key == "decoder.conv1.bias": new_state_dict["decoder.conv_in.bias"] = value - # Handle encoder downsamples elif key.startswith("encoder.downsamples."): - # Convert to down_blocks new_key = key.replace("encoder.downsamples.", "encoder.down_blocks.") - - # Convert residual block naming but keep the original structure if ".residual.0.gamma" in new_key: new_key = new_key.replace(".residual.0.gamma", ".norm1.gamma") elif ".residual.2.bias" in new_key: @@ -136,16 +136,11 @@ def convert_vae(vae_ckpt_path): new_key = new_key.replace(".shortcut.bias", ".conv_shortcut.bias") elif ".shortcut.weight" in new_key: new_key = new_key.replace(".shortcut.weight", ".conv_shortcut.weight") - new_state_dict[new_key] = value - - # Handle decoder upsamples elif key.startswith("decoder.upsamples."): - # Convert to up_blocks parts = key.split(".") block_idx = int(parts[2]) - # Group residual blocks if "residual" in key: if block_idx in [0, 1, 2]: new_block_idx = 0 @@ -160,11 +155,9 @@ def convert_vae(vae_ckpt_path): new_block_idx = 3 resnet_idx = block_idx - 12 else: - # Keep as is for other blocks new_state_dict[key] = value continue - # Convert residual block naming if ".residual.0.gamma" in key: new_key = f"decoder.up_blocks.{new_block_idx}.resnets.{resnet_idx}.norm1.gamma" elif ".residual.2.bias" in key: @@ -179,10 +172,8 @@ def convert_vae(vae_ckpt_path): new_key = f"decoder.up_blocks.{new_block_idx}.resnets.{resnet_idx}.conv2.weight" else: new_key = key - new_state_dict[new_key] = value - # Handle shortcut connections elif ".shortcut." in key: if block_idx == 4: new_key = key.replace(".shortcut.", ".resnets.0.conv_shortcut.") @@ -190,35 +181,22 @@ def convert_vae(vae_ckpt_path): else: new_key = key.replace("decoder.upsamples.", "decoder.up_blocks.") new_key = new_key.replace(".shortcut.", ".conv_shortcut.") - new_state_dict[new_key] = value - # Handle upsamplers elif ".resample." in key or ".time_conv." in key: if block_idx == 3: - new_key = key.replace( - f"decoder.upsamples.{block_idx}", - "decoder.up_blocks.0.upsamplers.0", - ) + new_key = key.replace(f"decoder.upsamples.{block_idx}", "decoder.up_blocks.0.upsamplers.0") elif block_idx == 7: - new_key = key.replace( - f"decoder.upsamples.{block_idx}", - "decoder.up_blocks.1.upsamplers.0", - ) + new_key = key.replace(f"decoder.upsamples.{block_idx}", "decoder.up_blocks.1.upsamplers.0") elif block_idx == 11: - new_key = key.replace( - f"decoder.upsamples.{block_idx}", - "decoder.up_blocks.2.upsamplers.0", - ) + new_key = key.replace(f"decoder.upsamples.{block_idx}", "decoder.up_blocks.2.upsamplers.0") else: new_key = key.replace("decoder.upsamples.", "decoder.up_blocks.") - new_state_dict[new_key] = value else: new_key = key.replace("decoder.upsamples.", "decoder.up_blocks.") new_state_dict[new_key] = value else: - # Keep other keys unchanged new_state_dict[key] = value with init_empty_weights(): @@ -227,32 +205,27 @@ def convert_vae(vae_ckpt_path): return vae -def get_transformer_config() -> Tuple[Dict[str, Any], ...]: - config = { - "diffusers_config": { - "hidden_size": 4096, - "in_channels": 16, - "num_attention_heads": 32, - "num_layers": 40, - "out_channels": 16, - "patch_size": [1, 2, 2], - "rope_dim_list": [16, 56, 56], - "text_dim": 4096, - "rope_type": "rope", - "theta": 10000, - }, - } - return config +TRANSFORMER_CONFIG = { + "hidden_size": 4096, + "in_channels": 16, + "num_attention_heads": 32, + "num_layers": 40, + "out_channels": 16, + "patch_size": [1, 2, 2], + "rope_dim_list": [16, 56, 56], + "text_dim": 4096, + "rope_type": "rope", + "theta": 10000, +} -def convert_transformer(ckpt_path: str): +def convert_transformer(ckpt_path: str, model_type: str = "edit"): checkpoint = torch.load(ckpt_path, weights_only=True) if "model" in checkpoint: original_state_dict = checkpoint["model"] else: original_state_dict = checkpoint - # Attention weights moved from block to block.attn submodule attn_suffixes = ( "img_attn_qkv.", "img_attn_q_norm.", @@ -268,21 +241,29 @@ def convert_transformer(ckpt_path: str): new_key = key if key.startswith("double_blocks."): for suffix in attn_suffixes: - # double_blocks.0.img_attn_qkv.weight -> double_blocks.0.attn.img_attn_qkv.weight if "." + suffix in key and ".attn." + suffix not in key: new_key = key.replace("." + suffix, ".attn." + suffix) break remapped[new_key] = value - config = get_transformer_config() + transformer_cls = ( + JoyImageEditPlusTransformer3DModel if model_type == "edit_plus" else JoyImageEditTransformer3DModel + ) with init_empty_weights(): - transformer = JoyImageEditTransformer3DModel(**config["diffusers_config"]) + transformer = transformer_cls(**TRANSFORMER_CONFIG) transformer.load_state_dict(remapped, strict=True, assign=True) return transformer def get_args(): - parser = argparse.ArgumentParser() + parser = argparse.ArgumentParser(description="Convert JoyImage Edit / Edit Plus checkpoints to diffusers format") + parser.add_argument( + "--model_type", + type=str, + choices=["edit", "edit_plus"], + default="edit", + help="Model type: 'edit' for JoyImage-Edit, 'edit_plus' for JoyImage-Edit-Plus", + ) parser.add_argument( "--transformer_ckpt_path", type=str, @@ -299,13 +280,7 @@ def get_args(): "--text_encoder_path", type=str, default=None, - help="Path to original llama checkpoint", - ) - parser.add_argument( - "--tokenizer_path", - type=str, - default=None, - help="Path to original llama tokenizer", + help="Path to Qwen3-VL text encoder (e.g. Qwen/Qwen3-VL-8B-Instruct)", ) parser.add_argument("--save_pipeline", action="store_true") parser.add_argument( @@ -314,8 +289,8 @@ def get_args(): required=True, help="Path where converted model should be saved", ) - parser.add_argument("--dtype", default="bf16", help="Torch dtype to save the transformer in.") - parser.add_argument("--flow_shift", type=float, default=7.0) + parser.add_argument("--dtype", default="bf16", help="Torch dtype (fp32, fp16, bf16)") + parser.add_argument("--flow_shift", type=float, default=1.5) return parser.parse_args() @@ -324,6 +299,7 @@ def get_args(): "fp16": torch.float16, "bf16": torch.bfloat16, } + if __name__ == "__main__": args = get_args() transformer = None @@ -333,34 +309,47 @@ def get_args(): if args.save_pipeline: assert args.transformer_ckpt_path is not None and args.vae_ckpt_path is not None assert args.text_encoder_path is not None - # assert args.tokenizer_path is not None + if args.transformer_ckpt_path is not None: - transformer = convert_transformer(args.transformer_ckpt_path) + transformer = convert_transformer(args.transformer_ckpt_path, model_type=args.model_type) transformer = transformer.to(dtype=dtype) if not args.save_pipeline: transformer.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB") + if args.vae_ckpt_path is not None: vae = convert_vae(args.vae_ckpt_path) vae = vae.to(dtype=dtype) if not args.save_pipeline: vae.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB") + if args.save_pipeline: processor = AutoProcessor.from_pretrained(args.text_encoder_path) text_encoder = Qwen3VLForConditionalGeneration.from_pretrained( args.text_encoder_path, torch_dtype=torch.bfloat16 ).to("cuda") tokenizer = AutoTokenizer.from_pretrained(args.text_encoder_path) - flow_shift = 1.5 - scheduler = FlowMatchEulerDiscreteScheduler(num_train_timesteps=1000, shift=flow_shift) + scheduler = FlowMatchEulerDiscreteScheduler(num_train_timesteps=1000, shift=args.flow_shift) transformer = transformer.to("cuda") vae = vae.to("cuda") - pipe = JoyImageEditPipeline( - processor=processor, - transformer=transformer, - text_encoder=text_encoder, - tokenizer=tokenizer, - vae=vae, - scheduler=scheduler, - ).to("cuda") + + if args.model_type == "edit_plus": + pipe = JoyImageEditPlusPipeline( + processor=processor, + transformer=transformer, + text_encoder=text_encoder, + tokenizer=tokenizer, + vae=vae, + scheduler=scheduler, + ).to("cuda") + else: + pipe = JoyImageEditPipeline( + processor=processor, + transformer=transformer, + text_encoder=text_encoder, + tokenizer=tokenizer, + vae=vae, + scheduler=scheduler, + ).to("cuda") + pipe.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB") processor.save_pretrained(f"{args.output_path}/processor") diff --git a/setup.py b/setup.py index f12f8fcf00c7..955f6feb09aa 100644 --- a/setup.py +++ b/setup.py @@ -114,6 +114,9 @@ "torchsde", "note_seq", "librosa", + # librosa pulls in numba/llvmlite; pin to versions that support Python 3.10+ + "llvmlite>=0.40.0", + "numba>=0.57.0", "numpy", "parameterized", "peft>=0.17.0", @@ -228,6 +231,8 @@ def run(self): "Jinja2", "invisible-watermark", "librosa", + "llvmlite", + "numba", "parameterized", "protobuf", "pytest", @@ -277,7 +282,7 @@ def run(self): setup( name="diffusers", - version="0.39.0.dev0", # expected format is one of x.y.z.dev0, or x.y.z.rc1 or x.y.z (no to dashes, yes to dots) + version="0.40.0.dev0", # expected format is one of x.y.z.dev0, or x.y.z.rc1 or x.y.z (no to dashes, yes to dots) description="State-of-the-art diffusion in PyTorch and JAX.", long_description=open("README.md", "r", encoding="utf-8").read(), long_description_content_type="text/markdown", diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index 81b36e113df4..dcccf5cd2de3 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -1,4 +1,4 @@ -__version__ = "0.39.0.dev0" +__version__ = "0.40.0.dev0" from typing import TYPE_CHECKING @@ -274,6 +274,7 @@ "HunyuanVideoTransformer3DModel", "I2VGenXLUNet", "Ideogram4Transformer2DModel", + "JoyImageEditPlusTransformer3DModel", "JoyImageEditTransformer3DModel", "Kandinsky3UNet", "Kandinsky5Transformer3DModel", @@ -476,6 +477,8 @@ [ "AnimaAutoBlocks", "AnimaModularPipeline", + "Cosmos3OmniBlocks", + "Cosmos3OmniModularPipeline", "ErnieImageAutoBlocks", "ErnieImageModularPipeline", "Flux2AutoBlocks", @@ -630,6 +633,8 @@ "ImageTextPipelineOutput", "JoyImageEditPipeline", "JoyImageEditPipelineOutput", + "JoyImageEditPlusPipeline", + "JoyImageEditPlusPipelineOutput", "Kandinsky3Img2ImgPipeline", "Kandinsky3Pipeline", "Kandinsky5I2IPipeline", @@ -1142,6 +1147,7 @@ HunyuanVideoTransformer3DModel, I2VGenXLUNet, Ideogram4Transformer2DModel, + JoyImageEditPlusTransformer3DModel, JoyImageEditTransformer3DModel, Kandinsky3UNet, Kandinsky5Transformer3DModel, @@ -1323,6 +1329,8 @@ from .modular_pipelines import ( AnimaAutoBlocks, AnimaModularPipeline, + Cosmos3OmniBlocks, + Cosmos3OmniModularPipeline, ErnieImageAutoBlocks, ErnieImageModularPipeline, Flux2AutoBlocks, @@ -1473,6 +1481,8 @@ ImageTextPipelineOutput, JoyImageEditPipeline, JoyImageEditPipelineOutput, + JoyImageEditPlusPipeline, + JoyImageEditPlusPipelineOutput, Kandinsky3Img2ImgPipeline, Kandinsky3Pipeline, Kandinsky5I2IPipeline, diff --git a/src/diffusers/dependency_versions_table.py b/src/diffusers/dependency_versions_table.py index e58981d3d4ab..fcced60419d0 100644 --- a/src/diffusers/dependency_versions_table.py +++ b/src/diffusers/dependency_versions_table.py @@ -21,6 +21,8 @@ "torchsde": "torchsde", "note_seq": "note_seq", "librosa": "librosa", + "llvmlite": "llvmlite>=0.40.0", + "numba": "numba>=0.57.0", "numpy": "numpy", "parameterized": "parameterized", "peft": "peft>=0.17.0", diff --git a/src/diffusers/loaders/lora_base.py b/src/diffusers/loaders/lora_base.py index 5c4a9903926a..898eae0ca36b 100644 --- a/src/diffusers/loaders/lora_base.py +++ b/src/diffusers/loaders/lora_base.py @@ -536,7 +536,7 @@ def unload_lora_weights(self): def fuse_lora( self, - components: list[str] = [], + components: list[str] | None = None, lora_scale: float = 1.0, safe_fusing: bool = False, adapter_names: list[str] | None = None, @@ -569,6 +569,9 @@ def fuse_lora( pipeline.fuse_lora(lora_scale=0.7) ``` """ + if components is None: + components = [] + if "fuse_unet" in kwargs: depr_message = "Passing `fuse_unet` to `fuse_lora()` is deprecated and will be ignored. Please use the `components` argument and provide a list of the components whose LoRAs are to be fused. `fuse_unet` will be removed in a future version." deprecate( @@ -620,7 +623,7 @@ def fuse_lora( self._merged_adapters = self._merged_adapters | merged_adapter_names - def unfuse_lora(self, components: list[str] = [], **kwargs): + def unfuse_lora(self, components: list[str] | None = None, **kwargs): r""" Reverses the effect of [`pipe.fuse_lora()`](https://huggingface.co/docs/diffusers/main/en/api/loaders#diffusers.loaders.LoraBaseMixin.fuse_lora). @@ -634,6 +637,9 @@ def unfuse_lora(self, components: list[str] = [], **kwargs): Whether to unfuse the text encoder LoRA parameters. If the text encoder wasn't monkey-patched with the LoRA parameters then it won't have any effect. """ + if components is None: + components = [] + if "unfuse_unet" in kwargs: depr_message = "Passing `unfuse_unet` to `unfuse_lora()` is deprecated and will be ignored. Please use the `components` argument. `unfuse_unet` will be removed in a future version." deprecate( diff --git a/src/diffusers/loaders/unet.py b/src/diffusers/loaders/unet.py index b7c996574d87..116d7d664647 100644 --- a/src/diffusers/loaders/unet.py +++ b/src/diffusers/loaders/unet.py @@ -33,14 +33,8 @@ from ..models.model_loading_utils import load_model_dict_into_meta from ..models.modeling_utils import _LOW_CPU_MEM_USAGE_DEFAULT, load_state_dict from ..utils import ( - USE_PEFT_BACKEND, _get_model_file, - convert_unet_state_dict_to_peft, - deprecate, - get_adapter_name, - get_peft_kwargs, is_accelerate_available, - is_peft_version, is_torch_version, logging, ) @@ -68,11 +62,11 @@ class UNet2DConditionLoadersMixin: @validate_hf_hub_args def load_attn_procs(self, pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor], **kwargs): r""" - Load pretrained attention processor layers into [`UNet2DConditionModel`]. Attention processor layers have to be - defined in + Load pretrained Custom Diffusion attention processor layers into [`UNet2DConditionModel`]. Attention processor + layers have to be defined in [`attention_processor.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py) - and be a `torch.nn.Module` class. Currently supported: LoRA, Custom Diffusion. For LoRA, one must install - `peft`: `pip install -U peft`. + and be a `torch.nn.Module` class. To load LoRA layers, use [`~loaders.PeftAdapterMixin.load_lora_adapter`] + instead. Parameters: pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`): @@ -106,31 +100,20 @@ def load_attn_procs(self, pretrained_model_name_or_path_or_dict: str | dict[str, allowed by Git. subfolder (`str`, *optional*, defaults to `""`): The subfolder location of a model file within a larger model repository on the Hub or locally. - network_alphas (`dict[str, float]`): - The value of the network alpha used for stable learning and preventing underflow. This value has the - same meaning as the `--network_alpha` option in the kohya-ss trainer script. Refer to [this - link](https://github.com/darkstorm2150/sd-scripts/blob/main/docs/train_network_README-en.md#execute-learning). - adapter_name (`str`, *optional*, defaults to None): - Adapter name to be used for referencing the loaded adapter model. If not specified, it will use - `default_{i}` where i is the total number of adapters being loaded. weight_name (`str`, *optional*, defaults to None): Name of the serialized state dict file. - low_cpu_mem_usage (`bool`, *optional*): - Speed up model loading by only loading the pretrained LoRA weights and not initializing the random - weights. Example: ```py - from diffusers import AutoPipelineForText2Image import torch + from diffusers import DiffusionPipeline - pipeline = AutoPipelineForText2Image.from_pretrained( - "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16 + pipeline = DiffusionPipeline.from_pretrained( + "CompVis/stable-diffusion-v1-4", + torch_dtype=torch.float16, ).to("cuda") - pipeline.unet.load_attn_procs( - "jbilcke-hf/sdxl-cinematic-1", weight_name="pytorch_lora_weights.safetensors", adapter_name="cinematic" - ) + pipeline.unet.load_attn_procs("path-to-save-model", weight_name="pytorch_custom_diffusion_weights.bin") ``` """ from ..hooks.group_offloading import _maybe_remove_and_reapply_group_offloading @@ -144,17 +127,9 @@ def load_attn_procs(self, pretrained_model_name_or_path_or_dict: str | dict[str, subfolder = kwargs.pop("subfolder", None) weight_name = kwargs.pop("weight_name", None) use_safetensors = kwargs.pop("use_safetensors", None) - adapter_name = kwargs.pop("adapter_name", None) _pipeline = kwargs.pop("_pipeline", None) - network_alphas = kwargs.pop("network_alphas", None) - low_cpu_mem_usage = kwargs.pop("low_cpu_mem_usage", _LOW_CPU_MEM_USAGE_DEFAULT) allow_pickle = False - if low_cpu_mem_usage and is_peft_version("<=", "0.13.0"): - raise ValueError( - "`low_cpu_mem_usage=True` is not compatible with this `peft` version. Please update it with `pip install -U peft`." - ) - if use_safetensors is None: use_safetensors = True allow_pickle = True @@ -204,42 +179,24 @@ def load_attn_procs(self, pretrained_model_name_or_path_or_dict: str | dict[str, state_dict = pretrained_model_name_or_path_or_dict is_custom_diffusion = any("custom_diffusion" in k for k in state_dict.keys()) - is_lora = all(("lora" in k or k.endswith(".alpha")) for k in state_dict.keys()) - is_model_cpu_offload = False - is_sequential_cpu_offload = False - is_group_offload = False - - if is_lora: - deprecation_message = "Using the `load_attn_procs()` method has been deprecated and will be removed in a future version. Please use `load_lora_adapter()`." - deprecate("load_attn_procs", "0.40.0", deprecation_message) - - if is_custom_diffusion: - attn_processors = self._process_custom_diffusion(state_dict=state_dict) - elif is_lora: - is_model_cpu_offload, is_sequential_cpu_offload, is_group_offload = self._process_lora( - state_dict=state_dict, - unet_identifier_key=self.unet_name, - network_alphas=network_alphas, - adapter_name=adapter_name, - _pipeline=_pipeline, - low_cpu_mem_usage=low_cpu_mem_usage, - ) - else: + if not is_custom_diffusion: raise ValueError( f"{model_file} does not seem to be in the correct format expected by Custom Diffusion training." ) + attn_processors = self._process_custom_diffusion(state_dict=state_dict) + # 0 else state_dict - - if len(state_dict_to_be_used) > 0: - if adapter_name in getattr(self, "peft_config", {}): - raise ValueError( - f"Adapter name {adapter_name} already in use in the Unet - please select a new adapter name." - ) - - state_dict = convert_unet_state_dict_to_peft(state_dict_to_be_used) - - if network_alphas is not None: - # The alphas state dict have the same structure as Unet, thus we convert it to peft format using - # `convert_unet_state_dict_to_peft` method. - network_alphas = convert_unet_state_dict_to_peft(network_alphas) - - rank = {} - for key, val in state_dict.items(): - if "lora_B" in key: - rank[key] = val.shape[1] - - lora_config_kwargs = get_peft_kwargs(rank, network_alphas, state_dict, is_unet=True) - if "use_dora" in lora_config_kwargs: - if lora_config_kwargs["use_dora"]: - if is_peft_version("<", "0.9.0"): - raise ValueError( - "You need `peft` 0.9.0 at least to use DoRA-enabled LoRAs. Please upgrade your installation of `peft`." - ) - else: - if is_peft_version("<", "0.9.0"): - lora_config_kwargs.pop("use_dora") - - if "lora_bias" in lora_config_kwargs: - if lora_config_kwargs["lora_bias"]: - if is_peft_version("<=", "0.13.2"): - raise ValueError( - "You need `peft` 0.14.0 at least to use `bias` in LoRAs. Please upgrade your installation of `peft`." - ) - else: - if is_peft_version("<=", "0.13.2"): - lora_config_kwargs.pop("lora_bias") - - lora_config = LoraConfig(**lora_config_kwargs) - - # adapter_name - if adapter_name is None: - adapter_name = get_adapter_name(self) - - # In case the pipeline has been already offloaded to CPU - temporarily remove the hooks - # otherwise loading LoRA weights will lead to an error - is_model_cpu_offload, is_sequential_cpu_offload, is_group_offload = self._optionally_disable_offloading( - _pipeline - ) - peft_kwargs = {} - if is_peft_version(">=", "0.13.1"): - peft_kwargs["low_cpu_mem_usage"] = low_cpu_mem_usage - - inject_adapter_in_model(lora_config, self, adapter_name=adapter_name, **peft_kwargs) - incompatible_keys = set_peft_model_state_dict(self, state_dict, adapter_name, **peft_kwargs) - - warn_msg = "" - if incompatible_keys is not None: - # Check only for unexpected keys. - unexpected_keys = getattr(incompatible_keys, "unexpected_keys", None) - if unexpected_keys: - lora_unexpected_keys = [k for k in unexpected_keys if "lora_" in k and adapter_name in k] - if lora_unexpected_keys: - warn_msg = ( - f"Loading adapter weights from state_dict led to unexpected keys found in the model:" - f" {', '.join(lora_unexpected_keys)}. " - ) - - # Filter missing keys specific to the current adapter. - missing_keys = getattr(incompatible_keys, "missing_keys", None) - if missing_keys: - lora_missing_keys = [k for k in missing_keys if "lora_" in k and adapter_name in k] - if lora_missing_keys: - warn_msg += ( - f"Loading adapter weights from state_dict led to missing keys in the model:" - f" {', '.join(lora_missing_keys)}." - ) - - if warn_msg: - logger.warning(warn_msg) - - return is_model_cpu_offload, is_sequential_cpu_offload, is_group_offload - @classmethod # Copied from diffusers.loaders.lora_base.LoraBaseMixin._optionally_disable_offloading def _optionally_disable_offloading(cls, _pipeline): @@ -420,8 +260,9 @@ def save_attn_procs( **kwargs, ): r""" - Save attention processor layers to a directory so that it can be reloaded with the - [`~loaders.UNet2DConditionLoadersMixin.load_attn_procs`] method. + Save Custom Diffusion attention processor layers to a directory so that it can be reloaded with the + [`~loaders.UNet2DConditionLoadersMixin.load_attn_procs`] method. To save LoRA layers, use + [`~loaders.PeftAdapterMixin.save_lora_adapter`] instead. Arguments: save_directory (`str` or `os.PathLike`): @@ -468,27 +309,22 @@ def save_attn_procs( ) for (_, x) in self.attn_processors.items() ) - if is_custom_diffusion: - state_dict = self._get_custom_diffusion_state_dict() - if save_function is None and safe_serialization: - # safetensors does not support saving dicts with non-tensor values - empty_state_dict = {k: v for k, v in state_dict.items() if not isinstance(v, torch.Tensor)} - if len(empty_state_dict) > 0: - logger.warning( - f"Safetensors does not support saving dicts with non-tensor values. " - f"The following keys will be ignored: {empty_state_dict.keys()}" - ) - state_dict = {k: v for k, v in state_dict.items() if isinstance(v, torch.Tensor)} - else: - deprecation_message = "Using the `save_attn_procs()` method has been deprecated and will be removed in a future version. Please use `save_lora_adapter()`." - deprecate("save_attn_procs", "0.40.0", deprecation_message) - - if not USE_PEFT_BACKEND: - raise ValueError("PEFT backend is required for saving LoRAs using the `save_attn_procs()` method.") - - from peft.utils import get_peft_model_state_dict + if not is_custom_diffusion: + raise ValueError( + "`save_attn_procs()` only supports saving Custom Diffusion attention processors. Please use " + "`save_lora_adapter()` to save LoRA layers." + ) - state_dict = get_peft_model_state_dict(self) + state_dict = self._get_custom_diffusion_state_dict() + if save_function is None and safe_serialization: + # safetensors does not support saving dicts with non-tensor values + empty_state_dict = {k: v for k, v in state_dict.items() if not isinstance(v, torch.Tensor)} + if len(empty_state_dict) > 0: + logger.warning( + f"Safetensors does not support saving dicts with non-tensor values. " + f"The following keys will be ignored: {empty_state_dict.keys()}" + ) + state_dict = {k: v for k, v in state_dict.items() if isinstance(v, torch.Tensor)} if save_function is None: if safe_serialization: @@ -503,9 +339,9 @@ def save_function(weights, filename): if weight_name is None: if safe_serialization: - weight_name = CUSTOM_DIFFUSION_WEIGHT_NAME_SAFE if is_custom_diffusion else LORA_WEIGHT_NAME_SAFE + weight_name = CUSTOM_DIFFUSION_WEIGHT_NAME_SAFE else: - weight_name = CUSTOM_DIFFUSION_WEIGHT_NAME if is_custom_diffusion else LORA_WEIGHT_NAME + weight_name = CUSTOM_DIFFUSION_WEIGHT_NAME # Save the model save_path = Path(save_directory, weight_name).as_posix() diff --git a/src/diffusers/models/__init__.py b/src/diffusers/models/__init__.py index c4a2fdaa682a..1746db347d32 100755 --- a/src/diffusers/models/__init__.py +++ b/src/diffusers/models/__init__.py @@ -121,6 +121,7 @@ _import_structure["transformers.transformer_hunyuanimage"] = ["HunyuanImageTransformer2DModel"] _import_structure["transformers.transformer_ideogram4"] = ["Ideogram4Transformer2DModel"] _import_structure["transformers.transformer_joyimage"] = ["JoyImageEditTransformer3DModel"] + _import_structure["transformers.transformer_joyimage_edit_plus"] = ["JoyImageEditPlusTransformer3DModel"] _import_structure["transformers.transformer_kandinsky"] = ["Kandinsky5Transformer3DModel"] _import_structure["transformers.transformer_krea2"] = ["Krea2Transformer2DModel"] _import_structure["transformers.transformer_longcat_audio_dit"] = ["LongCatAudioDiTTransformer"] @@ -254,6 +255,7 @@ HunyuanVideoFramepackTransformer3DModel, HunyuanVideoTransformer3DModel, Ideogram4Transformer2DModel, + JoyImageEditPlusTransformer3DModel, JoyImageEditTransformer3DModel, Kandinsky5Transformer3DModel, Krea2Transformer2DModel, diff --git a/src/diffusers/models/attention_dispatch.py b/src/diffusers/models/attention_dispatch.py index 3c86403ca206..9414c151fd67 100644 --- a/src/diffusers/models/attention_dispatch.py +++ b/src/diffusers/models/attention_dispatch.py @@ -340,6 +340,8 @@ class _HubKernelConfig: AttentionBackendName._FLASH_3_VARLEN_HUB: _HubKernelConfig( repo_id="kernels-community/flash-attn3", function_attr="flash_attn_varlen_func", + wrapped_forward_attr="flash_attn_interface._flash_attn_forward", + wrapped_backward_attr="flash_attn_interface._flash_attn_backward", version=1, ), AttentionBackendName.FLASH_HUB: _HubKernelConfig( @@ -1612,6 +1614,194 @@ def _flash_attention_3_hub_backward_op( return grad_query, grad_key, grad_value +def _flash_attention_3_varlen_hub_forward_op( + ctx: torch.autograd.function.FunctionCtx, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_mask: torch.Tensor | None = None, + dropout_p: float = 0.0, + is_causal: bool = False, + scale: float | None = None, + enable_gqa: bool = False, + return_lse: bool = False, + _save_ctx: bool = True, + _parallel_config: "ParallelConfig" | None = None, + *, + window_size: tuple[int, int] = (-1, -1), + softcap: float = 0.0, + num_splits: int = 1, + pack_gqa: bool | None = None, + deterministic: bool = False, + sm_margin: int = 0, +): + if dropout_p != 0.0: + raise ValueError("`dropout_p` is not yet supported for flash-attn 3 varlen hub kernels.") + if enable_gqa: + raise ValueError("`enable_gqa` is not yet supported for flash-attn 3 varlen hub kernels.") + + config = _HUB_KERNELS_REGISTRY[AttentionBackendName._FLASH_3_VARLEN_HUB] + wrapped_forward_fn = config.wrapped_forward_fn + wrapped_backward_fn = config.wrapped_backward_fn + if wrapped_forward_fn is None or wrapped_backward_fn is None: + raise RuntimeError( + "Flash attention 3 varlen hub kernels must expose `flash_attn_interface._flash_attn_forward` and " + "`flash_attn_interface._flash_attn_backward` for context parallel execution." + ) + + if scale is None: + scale = query.shape[-1] ** (-0.5) + + batch_size, seq_len_q, num_heads, _ = query.shape + _, seq_len_kv, _, _ = key.shape + + if attn_mask is not None: + attn_mask = _normalize_attn_mask(attn_mask, batch_size, seq_len_kv) + (_, seqlens_k), (cu_seqlens_q, cu_seqlens_k), (_, max_seqlen_k) = ( + _prepare_for_flash_attn_or_sage_varlen_with_mask(batch_size, seq_len_q, attn_mask, query.device) + ) + indices_k = attn_mask.flatten().nonzero(as_tuple=False).flatten() + query_packed = query.flatten(0, 1) + key_packed = key.reshape(-1, *key.shape[2:])[indices_k] + value_packed = value.reshape(-1, *value.shape[2:])[indices_k] + max_seqlen_q = seq_len_q + else: + (_, seqlens_k), (cu_seqlens_q, cu_seqlens_k), (max_seqlen_q, max_seqlen_k) = ( + _prepare_for_flash_attn_or_sage_varlen_without_mask(batch_size, seq_len_q, seq_len_kv, query.device) + ) + query_packed = query.flatten(0, 1) + key_packed = key.flatten(0, 1) + value_packed = value.flatten(0, 1) + seqlens_k = None + + out_packed, softmax_lse, *_ = wrapped_forward_fn( + query_packed, + key_packed, + value_packed, + None, # k_new + None, # v_new + None, # qv + None, # out_ + cu_seqlens_q, + cu_seqlens_k, + None, # cu_seqlens_k_new + None, # seqused_q + None, # seqused_k + max_seqlen_q, + max_seqlen_k, + None, # page_table + None, # kv_batch_idx + None, # leftpad_k + None, # rotary_cos + None, # rotary_sin + None, # seqlens_rotary + None, # q_descale + None, # k_descale + None, # v_descale + scale, + causal=is_causal, + window_size_left=window_size[0], + window_size_right=window_size[1], + attention_chunk=0, + softcap=softcap, + rotary_interleaved=True, + scheduler_metadata=None, + num_splits=num_splits, + pack_gqa=pack_gqa, + sm_margin=sm_margin, + ) + + out = out_packed.view(batch_size, seq_len_q, *out_packed.shape[1:]) + + if _save_ctx: + ctx.save_for_backward( + query_packed, key_packed, value_packed, out_packed, softmax_lse, cu_seqlens_q, cu_seqlens_k + ) + ctx.seqlens_k = seqlens_k # None if unmasked + ctx.indices_k = indices_k if attn_mask is not None else None + ctx.max_seqlen_q = max_seqlen_q + ctx.max_seqlen_k = max_seqlen_k + ctx.batch_size = batch_size + ctx.seq_len_q = seq_len_q + ctx.seq_len_kv = seq_len_kv + ctx.num_heads = num_heads + ctx.scale = scale + ctx.is_causal = is_causal + ctx.window_size = window_size + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.sm_margin = sm_margin + + # softmax_lse in varlen mode: (num_heads, total_q) -> (batch_size, seq_len_q, num_heads) + lse_sp = softmax_lse.view(num_heads, batch_size, seq_len_q).permute(1, 2, 0).contiguous() + + return (out, lse_sp) if return_lse else out + + +def _flash_attention_3_varlen_hub_backward_op( + ctx: torch.autograd.function.FunctionCtx, + grad_out: torch.Tensor, + *args, + **kwargs, +): + config = _HUB_KERNELS_REGISTRY[AttentionBackendName._FLASH_3_VARLEN_HUB] + wrapped_backward_fn = config.wrapped_backward_fn + if wrapped_backward_fn is None: + raise RuntimeError( + "Flash attention 3 varlen hub kernels must expose `flash_attn_interface._flash_attn_backward` " + "for context parallel execution." + ) + + query_packed, key_packed, value_packed, out_packed, softmax_lse, cu_seqlens_q, cu_seqlens_k = ctx.saved_tensors + + grad_out_packed = grad_out.flatten(0, 1) + grad_query, grad_key, grad_value = ( + torch.empty_like(query_packed), + torch.empty_like(key_packed), + torch.empty_like(value_packed), + ) + + wrapped_backward_fn( + grad_out_packed, + query_packed, + key_packed, + value_packed, + out_packed, + softmax_lse, + cu_seqlens_q, + cu_seqlens_k, + None, + None, # seqused_q, seqused_k + ctx.max_seqlen_q, + ctx.max_seqlen_k, + grad_query, + grad_key, + grad_value, + ctx.scale, + ctx.is_causal, + ctx.window_size[0], + ctx.window_size[1], + ctx.softcap, + ctx.deterministic, + ctx.sm_margin, + ) + + grad_query = grad_query.view(ctx.batch_size, ctx.seq_len_q, *grad_query.shape[1:]) + + if ctx.seqlens_k is not None: + grad_key = _unpad_to_padded(grad_key, ctx.indices_k, ctx.batch_size, ctx.seq_len_kv) + grad_value = _unpad_to_padded(grad_value, ctx.indices_k, ctx.batch_size, ctx.seq_len_kv) + else: + grad_key = grad_key.view(ctx.batch_size, ctx.seq_len_kv, *grad_key.shape[1:]) + grad_value = grad_value.view(ctx.batch_size, ctx.seq_len_kv, *grad_value.shape[1:]) + + grad_query = grad_query[..., : grad_out.shape[-1]] + grad_key = grad_key[..., : grad_out.shape[-1]] + grad_value = grad_value[..., : grad_out.shape[-1]] + + return grad_query, grad_key, grad_value + + def _sage_attention_forward_op( ctx: torch.autograd.function.FunctionCtx, query: torch.Tensor, @@ -3007,7 +3197,7 @@ def _flash_attention_3_hub( @_AttentionBackendRegistry.register( AttentionBackendName._FLASH_3_VARLEN_HUB, constraints=[_check_device, _check_qkv_dtype_bf16_or_fp16, _check_shape], - supports_context_parallel=False, + supports_context_parallel=True, ) def _flash_attention_3_varlen_hub( query: torch.Tensor, @@ -3019,41 +3209,74 @@ def _flash_attention_3_varlen_hub( return_lse: bool = False, _parallel_config: "ParallelConfig" | None = None, ) -> torch.Tensor: + if _parallel_config is not None and _parallel_config.context_parallel_config.ring_degree > 1: + raise NotImplementedError("`ring_degree > 1` is not yet supported for the _FLASH_3_VARLEN_HUB backend.") + batch_size, seq_len_q, _, _ = query.shape _, seq_len_kv, _, _ = key.shape - if attn_mask is not None: - attn_mask = _normalize_attn_mask(attn_mask, batch_size, seq_len_kv) - - (_, seqlens_k), (cu_seqlens_q, cu_seqlens_k), (max_seqlen_q, max_seqlen_k) = ( - _prepare_for_flash_attn_or_sage_varlen( - batch_size, seq_len_q, seq_len_kv, attn_mask=attn_mask, device=query.device - ) - ) - - key_valid, value_valid = [], [] - for b in range(batch_size): - valid_len = seqlens_k[b] - key_valid.append(key[b, :valid_len]) - value_valid.append(value[b, :valid_len]) + if _parallel_config is None: + if attn_mask is not None: + attn_mask = _normalize_attn_mask(attn_mask, batch_size, seq_len_kv) + (_, _), (cu_seqlens_q, cu_seqlens_k), (max_seqlen_q, max_seqlen_k) = ( + _prepare_for_flash_attn_or_sage_varlen_with_mask(batch_size, seq_len_q, attn_mask, query.device) + ) + indices_k = attn_mask.flatten().nonzero(as_tuple=False).flatten() + key_packed = key.reshape(-1, *key.shape[2:])[indices_k] + value_packed = value.reshape(-1, *value.shape[2:])[indices_k] + else: + (_, _), (cu_seqlens_q, cu_seqlens_k), (max_seqlen_q, max_seqlen_k) = ( + _prepare_for_flash_attn_or_sage_varlen_without_mask(batch_size, seq_len_q, seq_len_kv, query.device) + ) + key_packed = key.flatten(0, 1) + value_packed = value.flatten(0, 1) - query_packed = query.flatten(0, 1) - key_packed = torch.cat(key_valid, dim=0) - value_packed = torch.cat(value_valid, dim=0) + query_packed = query.flatten(0, 1) - func = _HUB_KERNELS_REGISTRY[AttentionBackendName._FLASH_3_VARLEN_HUB].kernel_fn - out, lse, *_ = func( - q=query_packed, - k=key_packed, - v=value_packed, - cu_seqlens_q=cu_seqlens_q, - cu_seqlens_k=cu_seqlens_k, - max_seqlen_q=max_seqlen_q, - max_seqlen_k=max_seqlen_k, - softmax_scale=scale, - causal=is_causal, - ) - out = out.unflatten(0, (batch_size, -1)) + func = _HUB_KERNELS_REGISTRY[AttentionBackendName._FLASH_3_VARLEN_HUB].kernel_fn + result = func( + q=query_packed, + k=key_packed, + v=value_packed, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + softmax_scale=scale, + causal=is_causal, + ) + if isinstance(result, tuple): + out, lse, *_ = result + else: + out = result + lse = None + out = out.unflatten(0, (batch_size, -1)) + else: + forward_op = functools.partial( + _flash_attention_3_varlen_hub_forward_op, + window_size=(-1, -1), + softcap=0.0, + num_splits=1, + pack_gqa=None, + deterministic=False, + sm_margin=0, + ) + out = _templated_context_parallel_attention( + query, + key, + value, + attn_mask, + 0.0, + is_causal, + scale, + False, + return_lse, + forward_op=forward_op, + backward_op=_flash_attention_3_varlen_hub_backward_op, + _parallel_config=_parallel_config, + ) + if return_lse: + out, lse = out return (out, lse) if return_lse else out diff --git a/src/diffusers/models/autoencoders/autoencoder_vidtok.py b/src/diffusers/models/autoencoders/autoencoder_vidtok.py index 63aadb2dbc9c..296c7bd8d85a 100644 --- a/src/diffusers/models/autoencoders/autoencoder_vidtok.py +++ b/src/diffusers/models/autoencoders/autoencoder_vidtok.py @@ -1502,5 +1502,5 @@ def forward( dec = dec[:, :, :-time_padding, :, :] if not return_dict: - return dec + return (dec,) return DecoderOutput(sample=dec) diff --git a/src/diffusers/models/modeling_utils.py b/src/diffusers/models/modeling_utils.py index 41b0f689d9a4..650ab4c7fcda 100644 --- a/src/diffusers/models/modeling_utils.py +++ b/src/diffusers/models/modeling_utils.py @@ -1509,6 +1509,8 @@ def cuda(self, *args, **kwargs): def to(self, *args, **kwargs): from ..hooks.group_offloading import _is_group_offload_enabled + fp32_modules = self._keep_in_fp32_modules or [] + device_arg_or_kwarg_present = any(isinstance(arg, torch.device) for arg in args) or "device" in kwargs dtype_present_in_args = "dtype" in kwargs @@ -1528,6 +1530,11 @@ def to(self, *args, **kwargs): dtype_present_in_args = True break + if dtype_present_in_args and fp32_modules is not None: + logger.warning( + f"There are modules in {self.__class__.__name__} that should be kept in float32: {fp32_modules}. Casting directly with `to()` can lead to inconsistent results; set `torch_dtype` in `from_pretrained()` instead to keep these modules in float32." + ) + if getattr(self, "is_quantized", False): if dtype_present_in_args: raise ValueError( diff --git a/src/diffusers/models/transformers/__init__.py b/src/diffusers/models/transformers/__init__.py index 4ba9703b5fc0..21f5cb853643 100755 --- a/src/diffusers/models/transformers/__init__.py +++ b/src/diffusers/models/transformers/__init__.py @@ -42,6 +42,7 @@ from .transformer_hunyuanimage import HunyuanImageTransformer2DModel from .transformer_ideogram4 import Ideogram4Transformer2DModel from .transformer_joyimage import JoyImageEditTransformer3DModel + from .transformer_joyimage_edit_plus import JoyImageEditPlusTransformer3DModel from .transformer_kandinsky import Kandinsky5Transformer3DModel from .transformer_krea2 import Krea2Transformer2DModel from .transformer_longcat_audio_dit import LongCatAudioDiTTransformer diff --git a/src/diffusers/models/transformers/transformer_joyimage_edit_plus.py b/src/diffusers/models/transformers/transformer_joyimage_edit_plus.py new file mode 100644 index 000000000000..4a13845faad3 --- /dev/null +++ b/src/diffusers/models/transformers/transformer_joyimage_edit_plus.py @@ -0,0 +1,539 @@ +# Copyright 2025 The JoyImage Team and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import inspect +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from ...configuration_utils import ConfigMixin, register_to_config +from ...utils import logging +from ..attention import AttentionMixin, AttentionModuleMixin, FeedForward +from ..attention_dispatch import dispatch_attention_fn +from ..embeddings import PixArtAlphaTextProjection, TimestepEmbedding, Timesteps +from ..modeling_outputs import Transformer2DModelOutput +from ..modeling_utils import ModelMixin +from ..normalization import FP32LayerNorm + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +def _apply_rotary_emb_batched( + xq: torch.Tensor, + xk: torch.Tensor, + freqs_cis: tuple[torch.Tensor, torch.Tensor], +) -> tuple[torch.Tensor, torch.Tensor]: + """RoPE for batched [B, S, D] freqs.""" + cos, sin = freqs_cis[0].to(xq.device), freqs_cis[1].to(xq.device) + + # batched: [B, S, D] -> [B, S, 1, D] + cos = cos.unsqueeze(2) + sin = sin.unsqueeze(2) + + def _rotate_half(x): + 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) + + 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 + + +# Copied from diffusers.models.transformers.transformer_joyimage.JoyImageModulate with JoyImage->JoyImageEditPlus +class JoyImageEditPlusModulate(nn.Module): + """Wan-style learnable modulation table. + + Produces `factor` modulation vectors by adding the conditioning signal to a learnable parameter table. + """ + + def __init__(self, hidden_size: int, factor: int, dtype=None, device=None): + super().__init__() + self.factor = factor + self.modulate_table = nn.Parameter( + torch.zeros(1, factor, hidden_size, dtype=dtype, device=device) / hidden_size**0.5, + requires_grad=True, + ) + + def forward(self, x: torch.Tensor) -> list[torch.Tensor]: + if x.ndim != 3: + x = x.unsqueeze(1) + return [o.squeeze(1) for o in (self.modulate_table + x).chunk(self.factor, dim=1)] + + +class JoyImageEditPlusAttnProcessor: + """Attention processor that supports batched RoPE embeddings for edit-plus multi-image input.""" + + _attention_backend = None + _parallel_config = None + + def __call__( + self, + attn: "JoyImageEditPlusAttention", + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor = None, + image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + attention_mask: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + if encoder_hidden_states is None: + raise ValueError("JoyImageEditPlusAttnProcessor requires encoder_hidden_states") + + heads = attn.heads + + img_qkv = attn.img_attn_qkv(hidden_states) + img_query, img_key, img_value = img_qkv.chunk(3, dim=-1) + + txt_qkv = attn.txt_attn_qkv(encoder_hidden_states) + txt_query, txt_key, txt_value = txt_qkv.chunk(3, dim=-1) + + img_query = img_query.unflatten(-1, (heads, -1)) + img_key = img_key.unflatten(-1, (heads, -1)) + img_value = img_value.unflatten(-1, (heads, -1)) + + txt_query = txt_query.unflatten(-1, (heads, -1)) + txt_key = txt_key.unflatten(-1, (heads, -1)) + txt_value = txt_value.unflatten(-1, (heads, -1)) + + img_query = attn.img_attn_q_norm(img_query) + img_key = attn.img_attn_k_norm(img_key) + txt_query = attn.txt_attn_q_norm(txt_query) + txt_key = attn.txt_attn_k_norm(txt_key) + + if image_rotary_emb is not None: + img_query, img_key = _apply_rotary_emb_batched(img_query, img_key, image_rotary_emb) + + 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) + + joint_hidden_states = dispatch_attention_fn( + joint_query, + joint_key, + joint_value, + attn_mask=attention_mask, + dropout_p=0.0, + is_causal=False, + backend=self._attention_backend, + parallel_config=self._parallel_config, + ) + + 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] :, :] + + img_attn_output = attn.img_attn_proj(img_attn_output) + txt_attn_output = attn.txt_attn_proj(txt_attn_output) + + return img_attn_output, txt_attn_output + + +class JoyImageEditPlusAttention(nn.Module, AttentionModuleMixin): + """Joint attention module for JoyImage Edit Plus double-stream blocks.""" + + _default_processor_cls = JoyImageEditPlusAttnProcessor + _available_processors = [JoyImageEditPlusAttnProcessor] + _supports_qkv_fusion = False + + def __init__( + self, + dim: int, + num_attention_heads: int, + attention_head_dim: int, + eps: float = 1e-6, + processor=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) + + if processor is None: + processor = self._default_processor_cls() + self.set_processor(processor) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | None = None, + image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + attention_mask: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + attn_parameters = set(inspect.signature(self.processor.__call__).parameters.keys()) + kwargs = {} + if "attention_mask" in attn_parameters: + kwargs["attention_mask"] = attention_mask + return self.processor(self, hidden_states, encoder_hidden_states, image_rotary_emb, **kwargs) + + +class JoyImageEditPlusTransformerBlock(nn.Module): + """Double-stream transformer block for JoyImage Edit Plus.""" + + def __init__( + self, + dim: int, + num_attention_heads: int, + attention_head_dim: int, + mlp_width_ratio: float = 4.0, + eps: float = 1e-6, + ): + super().__init__() + + self.dim = dim + self.num_attention_heads = num_attention_heads + self.attention_head_dim = attention_head_dim + mlp_hidden_dim = int(dim * mlp_width_ratio) + + # image stream + self.img_mod = JoyImageEditPlusModulate(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, activation_fn="gelu-approximate") + + # text stream + self.txt_mod = JoyImageEditPlusModulate(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, activation_fn="gelu-approximate") + + # joint attention + self.attn = JoyImageEditPlusAttention(dim, num_attention_heads, attention_head_dim, eps=eps) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + temb: torch.Tensor, + image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + attention_mask: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + # modulation + ( + 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) + + # --- attention --- + 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) + + # --- FFN --- + 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) + img_ffn_output = self.img_mlp(img_ffn_input) + txt_ffn_output = self.txt_mlp(txt_ffn_input) + hidden_states = hidden_states + img_ffn_output * img_mod2_gate.unsqueeze(1) + encoder_hidden_states = encoder_hidden_states + txt_ffn_output * txt_mod2_gate.unsqueeze(1) + + return hidden_states, encoder_hidden_states + + +# Copied from diffusers.models.transformers.transformer_joyimage.JoyImageTimeTextImageEmbedding with JoyImage->JoyImageEditPlus +class JoyImageEditPlusTimeTextImageEmbedding(nn.Module): + def __init__( + self, + dim: int, + time_freq_dim: int, + time_proj_dim: int, + text_embed_dim: int, + ): + 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, act_fn="gelu_tanh") + + def forward( + self, + timestep: torch.Tensor, + encoder_hidden_states: 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 JoyImageEditPlusTransformer3DModel(ModelMixin, ConfigMixin, AttentionMixin): + r""" + JoyImage Edit Plus Transformer for multi-image editing. + + Uses a patchify+padding approach where each reference image and the target noise are independently patchified and + concatenated into a flat patch sequence. Supports variable-resolution reference images. + + Input format: `[B, max_patches, C, pt, ph, pw]` (6D padded patches). + + Args: + patch_size (`list`, defaults to `[1, 2, 2]`): + Patch size for patchifying the latent input along `(t, h, w)` dimensions. + in_channels (`int`, defaults to `16`): + The number of channels in the input latent. + out_channels (`int`, *optional*, defaults to `None`): + The number of channels in the output. If not specified, it defaults to `in_channels`. + hidden_size (`int`, defaults to `3072`): + The dimensionality of the hidden representations. + num_attention_heads (`int`, defaults to `24`): + The number of attention heads. + text_dim (`int`, defaults to `4096`): + The dimensionality of the text encoder output. + mlp_width_ratio (`float`, defaults to `4.0`): + The ratio of MLP hidden dimension to `hidden_size`. + num_layers (`int`, defaults to `20`): + The number of double-stream transformer blocks. + rope_dim_list (`list[int]`, defaults to `[16, 56, 56]`): + The dimensions for 3D rotary positional embeddings along `(t, h, w)`. + rope_type (`str`, defaults to `"rope"`): + The type of rotary positional embedding. + theta (`int`, defaults to `256`): + The base frequency for rotary embeddings. + """ + + _skip_layerwise_casting_patterns = ["img_in", "condition_embedder", "norm"] + _no_split_modules = ["JoyImageEditPlusTransformerBlock"] + _supports_gradient_checkpointing = True + _keep_in_fp32_modules = [ + "time_embedder", + "norm1", + "norm2", + "norm_out", + ] + _repeated_blocks = ["JoyImageEditPlusTransformerBlock"] + + @register_to_config + def __init__( + self, + patch_size: list[int] = [1, 2, 2], + in_channels: int = 16, + out_channels: int | None = None, + hidden_size: int = 3072, + num_attention_heads: int = 24, + text_dim: int = 4096, + mlp_width_ratio: float = 4.0, + num_layers: int = 20, + rope_dim_list: list[int] = [16, 56, 56], + rope_type: str = "rope", + theta: int = 256, + ): + super().__init__() + + self.out_channels = out_channels or in_channels + + attention_head_dim = hidden_size // num_attention_heads + if hidden_size % num_attention_heads != 0: + raise ValueError( + f"hidden_size ({hidden_size}) must be divisible by num_attention_heads ({num_attention_heads})" + ) + + self.img_in = nn.Conv3d(in_channels, hidden_size, kernel_size=patch_size, stride=patch_size) + + self.condition_embedder = JoyImageEditPlusTimeTextImageEmbedding( + dim=hidden_size, + time_freq_dim=256, + time_proj_dim=hidden_size * 6, + text_embed_dim=text_dim, + ) + + self.double_blocks = nn.ModuleList( + [ + JoyImageEditPlusTransformerBlock( + dim=hidden_size, + num_attention_heads=num_attention_heads, + attention_head_dim=attention_head_dim, + mlp_width_ratio=mlp_width_ratio, + ) + for _ 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(patch_size)) + + self.gradient_checkpointing = False + + # Set batched-RoPE-aware attention processor on all blocks + for block in self.double_blocks: + block.attn.set_processor(JoyImageEditPlusAttnProcessor()) + + def _get_rotary_pos_embed_for_range( + self, + start: tuple[int, int, int], + stop: tuple[int, int, int], + ) -> tuple[torch.Tensor, torch.Tensor]: + """Generate 3D RoPE for a spatial range [start, stop).""" + head_dim = self.config.hidden_size // self.config.num_attention_heads + rope_dim_list = self.config.rope_dim_list + if rope_dim_list is None: + rope_dim_list = [head_dim // 3] * 3 + + grids = [] + for i in range(3): + grids.append(torch.arange(start[i], stop[i], dtype=torch.float32)) + + mesh = torch.stack(torch.meshgrid(*grids, indexing="ij"), dim=0) + + cos_parts, sin_parts = [], [] + for i, dim in enumerate(rope_dim_list): + pos = mesh[i].reshape(-1) + freqs = 1.0 / (self.config.theta ** (torch.arange(0, dim, 2, dtype=torch.float32)[: (dim // 2)] / dim)) + angles = torch.outer(pos, freqs) + cos_parts.append(angles.cos().repeat_interleave(2, dim=1)) + sin_parts.append(angles.sin().repeat_interleave(2, dim=1)) + + return torch.cat(cos_parts, dim=1), torch.cat(sin_parts, dim=1) + + def forward( + self, + hidden_states: torch.Tensor, + timestep: torch.Tensor, + encoder_hidden_states: torch.Tensor, + encoder_hidden_states_mask: torch.Tensor | None = None, + shape_list: list[list[tuple[int, int, int]]] = None, + return_dict: bool = True, + ) -> torch.Tensor | tuple: + """ + Args: + hidden_states: [B, max_patches, C, pt, ph, pw] - patchified latent input. + timestep: [B] - diffusion timestep. + encoder_hidden_states: [B, L, D] - text encoder outputs. + encoder_hidden_states_mask: [B, L] - attention mask for text tokens. + shape_list: Per-sample list of (t, h, w) tuples for each component (target + references). + return_dict: Whether to return a dict or tuple. + + Returns: + If `return_dict` is True, an [`~models.modeling_outputs.Transformer2DModelOutput`] is returned, otherwise a + `tuple` where the first element is the sample tensor. + """ + batch_size, max_num_patches, channels, pt, ph, pw = hidden_states.shape + device = hidden_states.device + + # 1. Condition embeddings + _, vec, txt = self.condition_embedder(timestep, encoder_hidden_states) + vec = vec.unflatten(1, (6, -1)) + + # 2. Patchify via Conv3d: flatten (B, N) -> apply conv -> reshape back + x = hidden_states.reshape(batch_size * max_num_patches, channels, pt, ph, pw) + x = self.img_in(x) # (B*N, D, 1, 1, 1) + img = x.reshape(batch_size, max_num_patches, -1) + + # 3. Build per-component RoPE with temporal offsets + sample_cos_list, sample_sin_list = [], [] + + for i in range(batch_size): + s_cos_parts, s_sin_parts = [], [] + current_t_offset = 0 + + for thw in shape_list[i]: + t, h, w = thw + start = (current_t_offset, 0, 0) + stop = (current_t_offset + t, h, w) + cos_emb, sin_emb = self._get_rotary_pos_embed_for_range(start, stop) + s_cos_parts.append(cos_emb) + s_sin_parts.append(sin_emb) + current_t_offset += t + + s_cos = torch.cat(s_cos_parts, dim=0).to(device) + s_sin = torch.cat(s_sin_parts, dim=0).to(device) + + actual_len = s_cos.shape[0] + pad_len = max_num_patches - actual_len + if pad_len > 0: + s_cos = F.pad(s_cos, (0, 0, 0, pad_len), value=1.0) + s_sin = F.pad(s_sin, (0, 0, 0, pad_len), value=0.0) + + sample_cos_list.append(s_cos) + sample_sin_list.append(s_sin) + + vis_freqs = (torch.stack(sample_cos_list), torch.stack(sample_sin_list)) + + # 4. Build attention mask: [B, 1, 1, img_seq + txt_seq] + attention_mask = None + if encoder_hidden_states_mask is not None: + img_mask = torch.zeros(batch_size, max_num_patches, device=device, dtype=encoder_hidden_states_mask.dtype) + for i in range(batch_size): + actual_len = sum(t * h * w for t, h, w in shape_list[i]) + img_mask[i, :actual_len] = 1.0 + full_mask = torch.cat([img_mask, encoder_hidden_states_mask], dim=1) + attention_mask = full_mask.unsqueeze(1).unsqueeze(1).bool() + + # 5. Run double blocks + for block in self.double_blocks: + if torch.is_grad_enabled() and self.gradient_checkpointing: + img, txt = self._gradient_checkpointing_func(block, img, txt, vec, vis_freqs, attention_mask) + else: + img, txt = block( + hidden_states=img, + encoder_hidden_states=txt, + temb=vec, + image_rotary_emb=vis_freqs, + attention_mask=attention_mask, + ) + + # 6. Output projection + reshape to 6D patches + img = self.proj_out(self.norm_out(img)) + img = img.reshape(batch_size, max_num_patches, pt, ph, pw, self.out_channels).permute( + 0, 1, 5, 2, 3, 4 + ) # -> [B, N, C, pt, ph, pw] + + if not return_dict: + return (img,) + return Transformer2DModelOutput(sample=img) diff --git a/src/diffusers/modular_pipelines/__init__.py b/src/diffusers/modular_pipelines/__init__.py index 4b36994aef07..25db2ef3bee2 100644 --- a/src/diffusers/modular_pipelines/__init__.py +++ b/src/diffusers/modular_pipelines/__init__.py @@ -97,6 +97,10 @@ "AnimaAutoBlocks", "AnimaModularPipeline", ] + _import_structure["cosmos"] = [ + "Cosmos3OmniBlocks", + "Cosmos3OmniModularPipeline", + ] _import_structure["ernie_image"] = [ "ErnieImageAutoBlocks", "ErnieImageModularPipeline", @@ -124,6 +128,7 @@ else: from .anima import AnimaAutoBlocks, AnimaModularPipeline from .components_manager import ComponentsManager + from .cosmos import Cosmos3OmniBlocks, Cosmos3OmniModularPipeline from .ernie_image import ErnieImageAutoBlocks, ErnieImageModularPipeline from .flux import FluxAutoBlocks, FluxKontextAutoBlocks, FluxKontextModularPipeline, FluxModularPipeline from .flux2 import ( diff --git a/src/diffusers/modular_pipelines/cosmos/__init__.py b/src/diffusers/modular_pipelines/cosmos/__init__.py new file mode 100644 index 000000000000..20150b299893 --- /dev/null +++ b/src/diffusers/modular_pipelines/cosmos/__init__.py @@ -0,0 +1,47 @@ +from typing import TYPE_CHECKING + +from ...utils import ( + DIFFUSERS_SLOW_IMPORT, + OptionalDependencyNotAvailable, + _LazyModule, + get_objects_from_module, + is_torch_available, + is_transformers_available, +) + + +_dummy_objects = {} +_import_structure = {} + +try: + if not (is_transformers_available() and is_torch_available()): + raise OptionalDependencyNotAvailable() +except OptionalDependencyNotAvailable: + from ...utils import dummy_torch_and_transformers_objects # noqa F403 + + _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects)) +else: + _import_structure["modular_blocks_cosmos3"] = ["Cosmos3OmniBlocks"] + _import_structure["modular_pipeline"] = ["Cosmos3OmniModularPipeline"] + +if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: + try: + if not (is_transformers_available() and is_torch_available()): + raise OptionalDependencyNotAvailable() + except OptionalDependencyNotAvailable: + from ...utils.dummy_torch_and_transformers_objects import * # noqa F403 + else: + from .modular_blocks_cosmos3 import Cosmos3OmniBlocks + from .modular_pipeline import Cosmos3OmniModularPipeline +else: + import sys + + sys.modules[__name__] = _LazyModule( + __name__, + globals()["__file__"], + _import_structure, + module_spec=__spec__, + ) + + for name, value in _dummy_objects.items(): + setattr(sys.modules[__name__], name, value) diff --git a/src/diffusers/modular_pipelines/cosmos/after_decode.py b/src/diffusers/modular_pipelines/cosmos/after_decode.py new file mode 100644 index 000000000000..7f8dd903d615 --- /dev/null +++ b/src/diffusers/modular_pipelines/cosmos/after_decode.py @@ -0,0 +1,113 @@ +import torch + +from ...utils import encode_video, export_to_video +from ..modular_pipeline import ModularPipelineBlocks, PipelineState +from ..modular_pipeline_utils import InputParam, OutputParam +from .modular_pipeline import Cosmos3OmniModularPipeline + + +class Cosmos3ActionOutputStep(ModularPipelineBlocks): + model_name = "cosmos3-omni" + + @property + def description(self) -> str: + return "Post-processes action latents into action outputs." + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + name="action_latents", + type_hint=torch.Tensor, + default=None, + description="Denoised action latents.", + ), + InputParam( + name="action_mode", type_hint=str, default=None, description="Requested action-generation mode." + ), + InputParam( + name="raw_action_dim_resolved", + type_hint=int, + default=None, + description="Unpadded action-vector dimension.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [OutputParam("action", type_hint=list[torch.Tensor], description="Generated action vectors.")] + + @torch.no_grad() + def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + action_output = None + if block_state.action_mode in {"inverse_dynamics", "policy"} and block_state.action_latents is not None: + action_output = block_state.action_latents + if block_state.raw_action_dim_resolved is not None: + action_output = action_output[:, : block_state.raw_action_dim_resolved] + action_output = [action_output.detach().cpu()] + block_state.action = action_output + self.set_block_state(state, block_state) + return components, state + + +class Cosmos3ExportStep(ModularPipelineBlocks): + model_name = "cosmos3-omni" + + @property + def description(self) -> str: + return ( + "Optional export block that writes decoded outputs to disk. Writes `videos` to `output_path` via " + "`export_to_video`, or muxes `videos` with `sound` via `encode_video` when a waveform is present. " + "Not wired into the default blocks; add it explicitly when you want the pipeline to produce a file." + ) + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam(name="videos", required=True, description="Generated video frames to export."), + InputParam( + name="output_path", + type_hint=str, + required=True, + description="Destination path for the exported video.", + ), + InputParam(name="fps", type_hint=float, default=24.0, description="Frame rate of the exported video."), + InputParam( + name="sound", + type_hint=torch.Tensor, + default=None, + description="Generated waveform to mux into the video.", + ), + InputParam( + name="sampling_rate", + type_hint=int, + default=None, + description="Sample rate of the generated waveform in Hz.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [OutputParam("output_path", type_hint=str, description="Path of the exported video file.")] + + @torch.no_grad() + def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + output_path = str(block_state.output_path) + fps = int(round(block_state.fps)) + if block_state.sound is not None: + if block_state.sampling_rate is None: + raise ValueError("`sampling_rate` is required to export a video with sound.") + encode_video( + block_state.videos, + fps=fps, + audio=block_state.sound, + audio_sample_rate=int(block_state.sampling_rate), + output_path=output_path, + ) + else: + export_to_video(block_state.videos, output_path, fps=fps, macro_block_size=1) + block_state.output_path = output_path + self.set_block_state(state, block_state) + return components, state diff --git a/src/diffusers/modular_pipelines/cosmos/before_denoise.py b/src/diffusers/modular_pipelines/cosmos/before_denoise.py new file mode 100644 index 000000000000..534365b4ef89 --- /dev/null +++ b/src/diffusers/modular_pipelines/cosmos/before_denoise.py @@ -0,0 +1,1354 @@ +import copy +import math + +import torch + +from ...configuration_utils import FrozenDict +from ...models.transformers.transformer_cosmos3 import Cosmos3OmniTransformer +from ...pipelines.cosmos.pipeline_cosmos3_omni import _EMBODIMENT_TO_DOMAIN_ID, CosmosActionCondition +from ...schedulers import UniPCMultistepScheduler +from ...utils.torch_utils import randn_tensor +from ...video_processor import VideoProcessor +from ..modular_pipeline import ModularPipelineBlocks, PipelineState +from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam +from .modular_pipeline import Cosmos3OmniModularPipeline + + +class Cosmos3PrepareTextSegmentsStep(ModularPipelineBlocks): + model_name = "cosmos3-omni" + + @property + def description(self) -> str: + return "Builds cond/uncond text segments before denoising." + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("transformer", Cosmos3OmniTransformer)] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam(name="cond_input_ids", required=True, description="Token IDs for the conditional prompt."), + InputParam(name="uncond_input_ids", required=True, description="Token IDs for the unconditional prompt."), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "cond_text_segment", + type_hint=dict, + kwargs_type="denoiser_input_fields", + description="Conditional text segment for the denoiser.", + ), + OutputParam( + "uncond_text_segment", + type_hint=dict, + kwargs_type="denoiser_input_fields", + description="Unconditional text segment for the denoiser.", + ), + ] + + @torch.no_grad() + def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + device = components._execution_device + block_state.cond_text_segment = components._prepare_text_segment(block_state.cond_input_ids, device=device) + block_state.uncond_text_segment = components._prepare_text_segment(block_state.uncond_input_ids, device=device) + self.set_block_state(state, block_state) + return components, state + + +class Cosmos3VisionPrepareLatentsStep(ModularPipelineBlocks): + model_name = "cosmos3-omni" + + @property + def description(self) -> str: + return "Prepares noisy vision latents and the vision conditioning mask." + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("transformer", Cosmos3OmniTransformer)] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + name="x0_tokens_vision", + type_hint=torch.Tensor, + default=None, + description="Vision latents encoded from the conditioning image or video.", + ), + InputParam( + name="vision_condition_frames", + type_hint=list[int], + default=None, + description="Latent-frame indexes fixed by visual conditioning.", + ), + InputParam(name="num_frames", type_hint=int, required=True, description="Number of frames to generate."), + InputParam( + name="height", type_hint=int, required=True, description="Height of the generated video in pixels." + ), + InputParam( + name="width", type_hint=int, required=True, description="Width of the generated video in pixels." + ), + InputParam(name="fps", type_hint=float, default=24.0, description="Frame rate of the generated video."), + InputParam( + name="latents", + type_hint=torch.Tensor, + default=None, + description="Pre-generated noisy vision latents.", + ), + InputParam.template("generator"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam("latents", type_hint=torch.Tensor, description="Noisy vision latents for denoising."), + OutputParam("fps_vision", type_hint=float, description="Frame rate used to pack vision latents."), + OutputParam( + "vision_condition_mask", + type_hint=torch.Tensor, + kwargs_type="denoiser_input_fields", + description="Mask marking conditioned vision latent frames.", + ), + OutputParam( + "vision_condition_indexes_for_pack", + type_hint=list[int], + description="Indexes of conditioned vision latent frames.", + ), + ] + + @torch.no_grad() + def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + device = components._execution_device + dtype = components.transformer.dtype + + x0_tokens_vision = block_state.x0_tokens_vision + if x0_tokens_vision is None: + if block_state.num_frames < 1: + raise ValueError(f"num_frames must be >= 1, got {block_state.num_frames}.") + sf_spatial = components.vae_scale_factor_spatial + if block_state.height % sf_spatial != 0 or block_state.width % sf_spatial != 0: + raise ValueError( + f"height and width must be multiples of {sf_spatial}, got ({block_state.height}, {block_state.width})." + ) + latent_shape = ( + 1, + components.num_channels_latents, + (block_state.num_frames - 1) // components.vae_scale_factor_temporal + 1, + block_state.height // sf_spatial, + block_state.width // sf_spatial, + ) + x0_tokens_vision = torch.zeros(latent_shape, device=device, dtype=torch.float32) + else: + x0_tokens_vision = x0_tokens_vision.to(device=device, dtype=torch.float32) + + block_state.fps_vision = float(block_state.fps) + condition_frames = block_state.vision_condition_frames or [] + block_state.vision_condition_mask = torch.zeros((x0_tokens_vision.shape[2], 1, 1), device=device, dtype=dtype) + for frame_idx in condition_frames: + if 0 <= frame_idx < block_state.vision_condition_mask.shape[0]: + block_state.vision_condition_mask[frame_idx, 0, 0] = 1.0 + + if block_state.latents is None: + pure_noise = randn_tensor( + tuple(x0_tokens_vision.shape), generator=block_state.generator, device=device, dtype=dtype + ) + block_state.latents = ( + block_state.vision_condition_mask * x0_tokens_vision.to(device=device, dtype=dtype) + + (1.0 - block_state.vision_condition_mask) * pure_noise + ) + else: + block_state.latents = block_state.latents.to(device=device, dtype=dtype) + + vision_condition_indexes = torch.nonzero( + block_state.vision_condition_mask[:, 0, 0] > 0, as_tuple=False + ).flatten() + block_state.vision_condition_indexes_for_pack = [int(idx.item()) for idx in vision_condition_indexes] + + self.set_block_state(state, block_state) + return components, state + + +class Cosmos3SoundPrepareLatentsStep(ModularPipelineBlocks): + model_name = "cosmos3-omni" + + @property + def description(self) -> str: + return "Prepares noisy sound latents and the sound conditioning mask." + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("transformer", Cosmos3OmniTransformer), + ComponentSpec("scheduler", UniPCMultistepScheduler), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam(name="num_frames", type_hint=int, required=True, description="Number of frames to generate."), + InputParam(name="fps", type_hint=float, default=24.0, description="Frame rate of the generated video."), + InputParam( + name="sound_latents", + type_hint=torch.Tensor, + default=None, + description="Pre-generated noisy sound latents.", + ), + InputParam.template("generator"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam("sound_latents", type_hint=torch.Tensor, description="Noisy sound latents for denoising."), + OutputParam("fps_sound", type_hint=float, description="Frame rate of the sound latent sequence."), + OutputParam( + "sound_condition_mask", + type_hint=torch.Tensor, + kwargs_type="denoiser_input_fields", + description="Mask marking conditioned sound latent frames.", + ), + OutputParam("sound_scheduler", description="Scheduler used to update sound latents."), + ] + + @torch.no_grad() + def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + device = components._execution_device + dtype = components.transformer.dtype + + if not components.transformer.config.sound_gen: + raise ValueError("Sound generation requires a transformer trained with sound_gen=True.") + + sound_dim = components.transformer.config.sound_dim + block_state.fps_sound = float(components.transformer.config.sound_latent_fps) + n_audio_samples = int(block_state.num_frames / block_state.fps * components.sound_sampling_rate) + hop_size = components.sound_hop_size + t_sound = (n_audio_samples + hop_size - 1) // hop_size + x0_tokens_sound = torch.zeros(sound_dim, t_sound, device=device, dtype=dtype) + block_state.sound_condition_mask = torch.zeros((x0_tokens_sound.shape[1], 1), device=device, dtype=dtype) + + if block_state.sound_latents is None: + pure_noise = randn_tensor( + tuple(x0_tokens_sound.shape), generator=block_state.generator, device=device, dtype=dtype + ) + block_state.sound_latents = ( + block_state.sound_condition_mask.T * x0_tokens_sound + + (1.0 - block_state.sound_condition_mask.T) * pure_noise + ) + else: + block_state.sound_latents = block_state.sound_latents.to(device=device, dtype=dtype) + + block_state.sound_scheduler = copy.deepcopy(components.scheduler) + + self.set_block_state(state, block_state) + return components, state + + +class Cosmos3ActionPrepareLatentsStep(ModularPipelineBlocks): + model_name = "cosmos3-omni" + + @property + def description(self) -> str: + return "Prepares noisy action latents and the action conditioning mask." + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("transformer", Cosmos3OmniTransformer), + ComponentSpec("scheduler", UniPCMultistepScheduler), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + name="action", + type_hint=CosmosActionCondition, + required=True, + description="Action-conditioning metadata.", + ), + InputParam( + name="action_condition_frame_indexes", + type_hint=list[int], + default=None, + description="Action-frame indexes fixed by action conditioning.", + ), + InputParam( + name="action_latents", + type_hint=torch.Tensor, + default=None, + description="Pre-generated noisy action latents.", + ), + InputParam.template("generator"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam("action_latents", type_hint=torch.Tensor, description="Noisy action latents for denoising."), + OutputParam( + "action_condition_mask", + type_hint=torch.Tensor, + kwargs_type="denoiser_input_fields", + description="Mask marking conditioned action latent frames.", + ), + OutputParam( + "action_domain_ids", + type_hint=list[torch.Tensor], + kwargs_type="denoiser_input_fields", + description="Embodiment domain IDs for action conditioning.", + ), + OutputParam( + "raw_action_dim_resolved", + type_hint=int, + kwargs_type="denoiser_input_fields", + description="Unpadded action-vector dimension.", + ), + OutputParam("action_scheduler", description="Scheduler used to update action latents."), + ] + + @torch.no_grad() + def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + device = components._execution_device + dtype = components.transformer.dtype + action = block_state.action + + if not components.transformer.config.action_gen: + raise ValueError("action requires a transformer trained with action_gen=True.") + + block_state.raw_action_dim_resolved = int(action.raw_action_dim) if action.raw_action_dim is not None else None + if ( + block_state.raw_action_dim_resolved is not None + and block_state.raw_action_dim_resolved > components.transformer.config.action_dim + ): + raise ValueError( + f"raw_action_dim={block_state.raw_action_dim_resolved} exceeds the model action_dim=" + f"{components.transformer.config.action_dim}." + ) + + action_chunk_size = action.chunk_size + action_dim = components.transformer.action_dim + if action.mode == "forward_dynamics": + raw_actions = action.raw_actions + if raw_actions is None: + raise ValueError("action_mode='forward_dynamics' requires an action tensor.") + raw_actions = raw_actions.to(device=device, dtype=dtype) + if raw_actions.shape[-1] > action_dim: + raise ValueError( + f"Cosmos3 action dimension {raw_actions.shape[-1]} exceeds model action_dim={action_dim}." + ) + if raw_actions.shape[0] < action_chunk_size: + raw_actions = torch.cat( + [raw_actions, raw_actions[-1:].expand(action_chunk_size - raw_actions.shape[0], -1)], + dim=0, + ) + raw_actions = raw_actions[:action_chunk_size] + if raw_actions.shape[-1] < action_dim: + action_padding = torch.zeros( + raw_actions.shape[0], + action_dim - raw_actions.shape[-1], + dtype=raw_actions.dtype, + device=raw_actions.device, + ) + raw_actions = torch.cat([raw_actions, action_padding], dim=-1) + x0_tokens_action = raw_actions + else: + x0_tokens_action = torch.zeros(action_chunk_size, action_dim, device=device, dtype=dtype) + + if action.domain_name not in _EMBODIMENT_TO_DOMAIN_ID: + raise ValueError( + f"Unknown Cosmos3 action domain_name={action.domain_name!r}; expected one of {sorted(_EMBODIMENT_TO_DOMAIN_ID)}." + ) + block_state.action_domain_ids = [ + torch.tensor([_EMBODIMENT_TO_DOMAIN_ID[action.domain_name]], dtype=torch.long, device=device) + ] + condition_frames = block_state.action_condition_frame_indexes or [] + block_state.action_condition_mask = torch.zeros((x0_tokens_action.shape[0], 1), device=device, dtype=dtype) + for frame_idx in condition_frames: + if 0 <= frame_idx < block_state.action_condition_mask.shape[0]: + block_state.action_condition_mask[frame_idx, 0] = 1.0 + + if block_state.action_latents is None: + pure_noise = randn_tensor( + tuple(x0_tokens_action.shape), generator=block_state.generator, device=device, dtype=dtype + ) + block_state.action_latents = ( + block_state.action_condition_mask * x0_tokens_action + + (1.0 - block_state.action_condition_mask) * pure_noise + ) + if block_state.raw_action_dim_resolved is not None: + block_state.action_latents[:, block_state.raw_action_dim_resolved :] = 0 + else: + block_state.action_latents = block_state.action_latents.to(device=device, dtype=dtype) + + block_state.action_scheduler = copy.deepcopy(components.scheduler) + + self.set_block_state(state, block_state) + return components, state + + +class Cosmos3VisionPackSequenceStep(ModularPipelineBlocks): + model_name = "cosmos3-omni" + + @property + def description(self) -> str: + return "Builds separate cond/uncond vision sequence segments." + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("transformer", Cosmos3OmniTransformer)] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + name="cond_text_segment", type_hint=dict, required=True, description="Conditional text segment." + ), + InputParam( + name="uncond_text_segment", + type_hint=dict, + required=True, + description="Unconditional text segment.", + ), + InputParam( + name="latents", type_hint=torch.Tensor, required=True, description="Noisy vision latents to pack." + ), + InputParam( + name="fps_vision", + type_hint=float, + required=True, + description="Frame rate used to pack vision latents.", + ), + InputParam( + name="vision_condition_indexes_for_pack", + type_hint=list[int], + required=True, + description="Indexes of conditioned vision latent frames.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "cond_vision_segment", + type_hint=dict, + kwargs_type="denoiser_input_fields", + description="Conditional vision segment for the denoiser.", + ), + OutputParam( + "uncond_vision_segment", + type_hint=dict, + kwargs_type="denoiser_input_fields", + description="Unconditional vision segment for the denoiser.", + ), + ] + + @torch.no_grad() + def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + device = components._execution_device + has_image_condition = bool(block_state.vision_condition_indexes_for_pack) + + block_state.cond_vision_segment = components._prepare_vision_segment( + input_vision_tokens=block_state.latents, + has_image_condition=has_image_condition, + mrope_offset=block_state.cond_text_segment["vision_start_temporal_offset"], + vision_fps=block_state.fps_vision, + curr=block_state.cond_text_segment["und_len"], + device=device, + condition_frame_indexes=block_state.vision_condition_indexes_for_pack, + ) + block_state.uncond_vision_segment = components._prepare_vision_segment( + input_vision_tokens=block_state.latents, + has_image_condition=has_image_condition, + mrope_offset=block_state.uncond_text_segment["vision_start_temporal_offset"], + vision_fps=block_state.fps_vision, + curr=block_state.uncond_text_segment["und_len"], + device=device, + condition_frame_indexes=block_state.vision_condition_indexes_for_pack, + ) + + self.set_block_state(state, block_state) + return components, state + + +class Cosmos3SoundPackSequenceStep(ModularPipelineBlocks): + model_name = "cosmos3-omni" + + @property + def description(self) -> str: + return "Builds separate cond/uncond sound sequence segments." + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("transformer", Cosmos3OmniTransformer)] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + name="cond_text_segment", type_hint=dict, required=True, description="Conditional text segment." + ), + InputParam( + name="uncond_text_segment", + type_hint=dict, + required=True, + description="Unconditional text segment.", + ), + InputParam( + name="cond_sequence_length", + type_hint=int, + required=True, + description="Conditional multimodal sequence length.", + ), + InputParam( + name="uncond_sequence_length", + type_hint=int, + required=True, + description="Unconditional multimodal sequence length.", + ), + InputParam( + name="sound_latents", type_hint=torch.Tensor, required=True, description="Noisy sound latents to pack." + ), + InputParam( + name="fps_sound", + type_hint=float, + required=True, + description="Frame rate of the sound latent sequence.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "cond_sound_segment", + type_hint=dict, + kwargs_type="denoiser_input_fields", + description="Conditional sound segment for the denoiser.", + ), + OutputParam( + "uncond_sound_segment", + type_hint=dict, + kwargs_type="denoiser_input_fields", + description="Unconditional sound segment for the denoiser.", + ), + ] + + @torch.no_grad() + def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + device = components._execution_device + + block_state.cond_sound_segment = components._prepare_sound_segment( + input_sound_tokens=block_state.sound_latents, + mrope_offset=block_state.cond_text_segment["vision_start_temporal_offset"], + sound_fps=block_state.fps_sound, + curr=block_state.cond_sequence_length, + device=device, + ) + block_state.uncond_sound_segment = components._prepare_sound_segment( + input_sound_tokens=block_state.sound_latents, + mrope_offset=block_state.uncond_text_segment["vision_start_temporal_offset"], + sound_fps=block_state.fps_sound, + curr=block_state.uncond_sequence_length, + device=device, + ) + + self.set_block_state(state, block_state) + return components, state + + +class Cosmos3ActionPackSequenceStep(ModularPipelineBlocks): + model_name = "cosmos3-omni" + + @property + def description(self) -> str: + return "Builds separate cond/uncond action sequence segments." + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("transformer", Cosmos3OmniTransformer)] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + name="cond_text_segment", type_hint=dict, required=True, description="Conditional text segment." + ), + InputParam( + name="uncond_text_segment", + type_hint=dict, + required=True, + description="Unconditional text segment.", + ), + InputParam( + name="cond_sequence_length", + type_hint=int, + required=True, + description="Conditional multimodal sequence length.", + ), + InputParam( + name="uncond_sequence_length", + type_hint=int, + required=True, + description="Unconditional multimodal sequence length.", + ), + InputParam( + name="action_latents", + type_hint=torch.Tensor, + required=True, + description="Noisy action latents to pack.", + ), + InputParam( + name="action_condition_frame_indexes", + type_hint=list[int], + default=None, + description="Action-frame indexes fixed by action conditioning.", + ), + InputParam( + name="fps_vision", + type_hint=float, + required=True, + description="Frame rate used to pack vision latents.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "cond_action_segment", + type_hint=dict, + kwargs_type="denoiser_input_fields", + description="Conditional action segment for the denoiser.", + ), + OutputParam( + "uncond_action_segment", + type_hint=dict, + kwargs_type="denoiser_input_fields", + description="Unconditional action segment for the denoiser.", + ), + ] + + @torch.no_grad() + def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + device = components._execution_device + + block_state.cond_action_segment = components._prepare_action_segment( + input_action_tokens=block_state.action_latents, + condition_frame_indexes=block_state.action_condition_frame_indexes, + mrope_offset=block_state.cond_text_segment["vision_start_temporal_offset"], + action_fps=block_state.fps_vision, + curr=block_state.cond_sequence_length, + device=device, + ) + block_state.uncond_action_segment = components._prepare_action_segment( + input_action_tokens=block_state.action_latents, + condition_frame_indexes=block_state.action_condition_frame_indexes, + mrope_offset=block_state.uncond_text_segment["vision_start_temporal_offset"], + action_fps=block_state.fps_vision, + curr=block_state.uncond_sequence_length, + device=device, + ) + + self.set_block_state(state, block_state) + return components, state + + +class Cosmos3VisionDenoiseInputStep(ModularPipelineBlocks): + model_name = "cosmos3-omni" + + @property + def description(self) -> str: + return "Assembles text and vision sequence metadata for the denoising loop." + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + name="cond_text_segment", type_hint=dict, required=True, description="Conditional text segment." + ), + InputParam( + name="uncond_text_segment", + type_hint=dict, + required=True, + description="Unconditional text segment.", + ), + InputParam( + name="cond_vision_segment", type_hint=dict, required=True, description="Conditional vision segment." + ), + InputParam( + name="uncond_vision_segment", + type_hint=dict, + required=True, + description="Unconditional vision segment.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "cond_position_ids", + type_hint=torch.Tensor, + kwargs_type="denoiser_input_fields", + description="Conditional multimodal RoPE position IDs.", + ), + OutputParam( + "uncond_position_ids", + type_hint=torch.Tensor, + kwargs_type="denoiser_input_fields", + description="Unconditional multimodal RoPE position IDs.", + ), + OutputParam( + "cond_sequence_length", + type_hint=int, + kwargs_type="denoiser_input_fields", + description="Conditional multimodal sequence length.", + ), + OutputParam( + "uncond_sequence_length", + type_hint=int, + kwargs_type="denoiser_input_fields", + description="Unconditional multimodal sequence length.", + ), + ] + + @torch.no_grad() + def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + block_state.cond_position_ids = torch.cat( + [ + block_state.cond_text_segment["text_mrope_ids"], + block_state.cond_vision_segment["vision_mrope_ids"], + ], + dim=1, + ) + block_state.uncond_position_ids = torch.cat( + [ + block_state.uncond_text_segment["text_mrope_ids"], + block_state.uncond_vision_segment["vision_mrope_ids"], + ], + dim=1, + ) + block_state.cond_sequence_length = ( + block_state.cond_text_segment["und_len"] + block_state.cond_vision_segment["num_vision_tokens"] + ) + block_state.uncond_sequence_length = ( + block_state.uncond_text_segment["und_len"] + block_state.uncond_vision_segment["num_vision_tokens"] + ) + self.set_block_state(state, block_state) + return components, state + + +class Cosmos3SoundDenoiseInputStep(ModularPipelineBlocks): + model_name = "cosmos3-omni" + + @property + def description(self) -> str: + return "Appends sound sequence metadata to the denoising-loop inputs." + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + name="cond_position_ids", + type_hint=torch.Tensor, + required=True, + description="Conditional multimodal RoPE position IDs.", + ), + InputParam( + name="uncond_position_ids", + type_hint=torch.Tensor, + required=True, + description="Unconditional multimodal RoPE position IDs.", + ), + InputParam( + name="cond_sequence_length", + type_hint=int, + required=True, + description="Conditional multimodal sequence length.", + ), + InputParam( + name="uncond_sequence_length", + type_hint=int, + required=True, + description="Unconditional multimodal sequence length.", + ), + InputParam( + name="cond_sound_segment", type_hint=dict, required=True, description="Conditional sound segment." + ), + InputParam( + name="uncond_sound_segment", + type_hint=dict, + required=True, + description="Unconditional sound segment.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "cond_position_ids", + type_hint=torch.Tensor, + kwargs_type="denoiser_input_fields", + description="Conditional multimodal RoPE position IDs.", + ), + OutputParam( + "uncond_position_ids", + type_hint=torch.Tensor, + kwargs_type="denoiser_input_fields", + description="Unconditional multimodal RoPE position IDs.", + ), + OutputParam( + "cond_sequence_length", + type_hint=int, + kwargs_type="denoiser_input_fields", + description="Conditional multimodal sequence length.", + ), + OutputParam( + "uncond_sequence_length", + type_hint=int, + kwargs_type="denoiser_input_fields", + description="Unconditional multimodal sequence length.", + ), + ] + + @torch.no_grad() + def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + block_state.cond_position_ids = torch.cat( + [block_state.cond_position_ids, block_state.cond_sound_segment["sound_mrope_ids"]], dim=1 + ) + block_state.uncond_position_ids = torch.cat( + [block_state.uncond_position_ids, block_state.uncond_sound_segment["sound_mrope_ids"]], dim=1 + ) + block_state.cond_sequence_length += block_state.cond_sound_segment["sound_len"] + block_state.uncond_sequence_length += block_state.uncond_sound_segment["sound_len"] + self.set_block_state(state, block_state) + return components, state + + +class Cosmos3ActionDenoiseInputStep(ModularPipelineBlocks): + model_name = "cosmos3-omni" + + @property + def description(self) -> str: + return "Appends action sequence metadata to the denoising-loop inputs." + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + name="cond_position_ids", + type_hint=torch.Tensor, + required=True, + description="Conditional multimodal RoPE position IDs.", + ), + InputParam( + name="uncond_position_ids", + type_hint=torch.Tensor, + required=True, + description="Unconditional multimodal RoPE position IDs.", + ), + InputParam( + name="cond_sequence_length", + type_hint=int, + required=True, + description="Conditional multimodal sequence length.", + ), + InputParam( + name="uncond_sequence_length", + type_hint=int, + required=True, + description="Unconditional multimodal sequence length.", + ), + InputParam( + name="cond_action_segment", type_hint=dict, required=True, description="Conditional action segment." + ), + InputParam( + name="uncond_action_segment", + type_hint=dict, + required=True, + description="Unconditional action segment.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "cond_position_ids", + type_hint=torch.Tensor, + kwargs_type="denoiser_input_fields", + description="Conditional multimodal RoPE position IDs.", + ), + OutputParam( + "uncond_position_ids", + type_hint=torch.Tensor, + kwargs_type="denoiser_input_fields", + description="Unconditional multimodal RoPE position IDs.", + ), + OutputParam( + "cond_sequence_length", + type_hint=int, + kwargs_type="denoiser_input_fields", + description="Conditional multimodal sequence length.", + ), + OutputParam( + "uncond_sequence_length", + type_hint=int, + kwargs_type="denoiser_input_fields", + description="Unconditional multimodal sequence length.", + ), + ] + + @torch.no_grad() + def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + block_state.cond_position_ids = torch.cat( + [block_state.cond_position_ids, block_state.cond_action_segment["action_mrope_ids"]], dim=1 + ) + block_state.uncond_position_ids = torch.cat( + [block_state.uncond_position_ids, block_state.uncond_action_segment["action_mrope_ids"]], dim=1 + ) + block_state.cond_sequence_length += block_state.cond_action_segment["action_len"] + block_state.uncond_sequence_length += block_state.uncond_action_segment["action_len"] + self.set_block_state(state, block_state) + return components, state + + +class Cosmos3SetTimestepsStep(ModularPipelineBlocks): + model_name = "cosmos3-omni" + + @property + def description(self) -> str: + return "Initializes scheduler timesteps." + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("scheduler", UniPCMultistepScheduler)] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("num_inference_steps", required=True), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam("timesteps", type_hint=torch.Tensor, description="Scheduler timesteps for denoising."), + OutputParam("num_warmup_steps", type_hint=int, description="Number of scheduler warmup steps."), + ] + + @torch.no_grad() + def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + device = components._execution_device + components.scheduler.set_timesteps(block_state.num_inference_steps, device=device) + block_state.timesteps = components.scheduler.timesteps + block_state.num_warmup_steps = ( + len(block_state.timesteps) - block_state.num_inference_steps * components.scheduler.order + ) + self.set_block_state(state, block_state) + return components, state + + +class Cosmos3TransferSetupStep(ModularPipelineBlocks): + model_name = "cosmos3-omni" + + @property + def description(self) -> str: + return ( + "Preprocesses the transfer control videos and resolves the autoregressive chunk geometry " + "(total_frames / chunk_frames / num_chunks / stride). Chunk-invariant, so it runs once before the loop." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec( + "video_processor", + VideoProcessor, + config=FrozenDict({"vae_scale_factor": 16, "resample": "bilinear"}), + default_creation_method="from_config", + ), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + name="control_videos", + type_hint=dict, + required=True, + description="Mapping of hint name (edge/blur/depth/seg/wsm) to the control video for that modality.", + ), + InputParam( + name="height", type_hint=int, default=None, description="Height of the generated video in pixels." + ), + InputParam( + name="width", type_hint=int, default=None, description="Width of the generated video in pixels." + ), + InputParam( + name="num_frames", + type_hint=int, + default=None, + description="Optional cap on the number of output frames (defaults to the control video length).", + ), + InputParam( + name="num_video_frames_per_chunk", + type_hint=int, + default=None, + description="Number of pixel frames generated per autoregressive chunk.", + ), + InputParam( + name="num_conditional_frames", + type_hint=int, + default=1, + description="Number of frames each chunk reuses from the previous chunk's tail.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam("height", type_hint=int, description="Resolved output height in pixels."), + OutputParam("width", type_hint=int, description="Resolved output width in pixels."), + OutputParam( + "control_frames", + type_hint=dict, + description="Preprocessed, time-padded control maps in canonical hint order.", + ), + OutputParam("total_frames", type_hint=int, description="Total number of output frames to generate."), + OutputParam("chunk_frames", type_hint=int, description="Number of pixel frames per autoregressive chunk."), + OutputParam("num_chunks", type_hint=int, description="Number of autoregressive chunks."), + OutputParam("stride", type_hint=int, description="Frame stride between consecutive chunks."), + ] + + @torch.no_grad() + def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + device = components._execution_device + dtype = components.transformer.dtype + + if block_state.height is None: + block_state.height = 720 + if block_state.width is None: + block_state.width = 1280 + + # Canonical hint order used both to validate and to order the preprocessed control maps. + hint_order = ["edge", "blur", "depth", "seg", "wsm"] + control_videos = block_state.control_videos + if not isinstance(control_videos, dict) or not control_videos: + raise ValueError("`control_videos` must be a non-empty dict mapping hint name -> control video.") + unknown = [k for k in control_videos if k not in hint_order] + if unknown: + raise ValueError(f"`control_videos` has unknown hint(s) {unknown}; expected keys from {hint_order}.") + if any(v is None for v in control_videos.values()): + raise ValueError("`control_videos` entries must be loaded videos, not None.") + + tcf = components.vae_scale_factor_temporal + sf = components.vae_scale_factor_spatial + if block_state.height % sf != 0 or block_state.width % sf != 0: + raise ValueError( + f"`height` and `width` must be multiples of {sf}, got ({block_state.height}, {block_state.width})." + ) + + # Preprocess every control map to [1, 3, T, H, W] in [-1, 1] at target geometry, in canonical hint order. + # The dict preserves this order, so downstream blocks just iterate control_frames (no separate hint_keys). + hint_keys = [k for k in hint_order if k in control_videos] + control_frames = { + key: components.video_processor.preprocess_video( + control_videos[key], height=block_state.height, width=block_state.width + ).to(device=device, dtype=dtype) + for key in hint_keys + } + + # Output frame count / chunking come from the (first) control video, optionally capped by num_frames. + total_frames = next(iter(control_frames.values())).shape[2] + if block_state.num_frames is not None: + total_frames = min(total_frames, block_state.num_frames) + total_frames = max(1, total_frames) + + per_chunk = ( + block_state.num_video_frames_per_chunk + if block_state.num_video_frames_per_chunk is not None + else total_frames + ) + chunk_frames = 1 if total_frames == 1 else per_chunk + chunk_frames = math.ceil((chunk_frames - 1) / tcf) * tcf + 1 + + if total_frames <= chunk_frames: + num_chunks, stride = 1, chunk_frames + else: + stride = chunk_frames - block_state.num_conditional_frames + if stride <= 0: + raise ValueError("`num_conditional_frames` must be smaller than `num_video_frames_per_chunk`.") + remaining = total_frames - chunk_frames + num_chunks = 1 + (remaining // stride + (1 if remaining % stride else 0)) + + # Reflect-pad each control map along time up to `padded` (repeat the last frame once the clip is too short to + # keep reflecting). No truncation here; per-chunk slicing happens later. + padded = max(total_frames, chunk_frames) + control_frames_padded = {} + for key, frames in control_frames.items(): + while frames.shape[2] < padded: + pad_len = min(frames.shape[2] - 1, padded - frames.shape[2]) + if pad_len <= 0: + pad_frame = frames[:, :, -1:].repeat(1, 1, padded - frames.shape[2], 1, 1) + frames = torch.cat([frames, pad_frame], dim=2) + break + frames = torch.cat([frames, frames.flip(dims=[2])[:, :, :pad_len]], dim=2) + control_frames_padded[key] = frames + block_state.control_frames = control_frames_padded + block_state.total_frames = total_frames + block_state.chunk_frames = chunk_frames + block_state.num_chunks = num_chunks + block_state.stride = stride + + self.set_block_state(state, block_state) + return components, state + + +class Cosmos3TransferPrepareLatentsStep(ModularPipelineBlocks): + model_name = "cosmos3-omni" + + @property + def description(self) -> str: + return ( + "Per-chunk transfer latent prep: takes the clean target latents encoded by " + "Cosmos3TransferChunkVaeEncoderStep and builds the noisy target latents, velocity mask, condition latents " + "and conditioned-frame indexes for this chunk." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("transformer", Cosmos3OmniTransformer)] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + name="x0_tokens_vision", + type_hint=torch.Tensor, + required=True, + description="Clean target vision latents encoded from the seeded target frames.", + ), + InputParam( + name="current_conditional_frames", + type_hint=int, + required=True, + description="Number of pixel frames used to seed this chunk's target.", + ), + InputParam.template("generator"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam("latents", type_hint=torch.Tensor, description="Noisy target latents for this chunk."), + OutputParam( + "velocity_mask", + type_hint=torch.Tensor, + description="Mask that zeroes the velocity on conditioned (clean) latent frames.", + ), + OutputParam( + "condition_latents", + type_hint=torch.Tensor, + description="Clean target latents on the conditioned frames (the autoregressive seed).", + ), + OutputParam( + "target_condition_indexes", + type_hint=list[int], + description="Latent-frame indexes fixed by the chunk's conditioning.", + ), + ] + + @torch.no_grad() + def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + device = components._execution_device + dtype = components.transformer.dtype + tcf = components.vae_scale_factor_temporal + + target_x0 = block_state.x0_tokens_vision.to(device=device) + current_conditional_frames = block_state.current_conditional_frames + + # Build the noisy target latents + conditioning mask from the clean target latents. + latent_t = target_x0.shape[2] + condition_mask = torch.zeros((latent_t, 1, 1), device=device, dtype=dtype) + latent_condition_frames = 0 + if current_conditional_frames > 0: + latent_condition_frames = (current_conditional_frames - 1) // tcf + 1 + condition_mask[:latent_condition_frames] = 1.0 + noise = randn_tensor(tuple(target_x0.shape), generator=block_state.generator, device=device, dtype=dtype) + block_state.latents = condition_mask * target_x0 + (1.0 - condition_mask) * noise + block_state.velocity_mask = 1.0 - condition_mask + block_state.condition_latents = condition_mask * target_x0 + block_state.target_condition_indexes = list(range(latent_condition_frames)) + + self.set_block_state(state, block_state) + return components, state + + +class Cosmos3TransferPackSequenceStep(ModularPipelineBlocks): + model_name = "cosmos3-omni" + + @property + def description(self) -> str: + return ( + "Pre-packs the three transfer CFG sequence variants: cond_full / uncond_full carry every control item, " + "the no-control branch drops them (only [text, target]) so the control axis can be amplified." + ) + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam(name="cond_text_segment", type_hint=dict, required=True, description="Conditional text segment."), + InputParam( + name="uncond_text_segment", type_hint=dict, required=True, description="Unconditional text segment." + ), + InputParam( + name="control_latents", + type_hint=list[torch.Tensor], + required=True, + description="Clean control latents for this chunk, one per hint in canonical order.", + ), + InputParam( + name="latents", type_hint=torch.Tensor, required=True, description="Noisy target latents for this chunk." + ), + InputParam( + name="target_condition_indexes", + type_hint=list[int], + required=True, + description="Latent-frame indexes fixed by the chunk's conditioning.", + ), + InputParam(name="fps", type_hint=float, default=24.0, description="Frame rate of the generated video."), + InputParam( + name="share_vision_temporal_positions", + type_hint=bool, + default=True, + description="Whether control and target items share vision temporal positions.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "cond_full_static", + type_hint=dict, + kwargs_type="denoiser_input_fields", + description="Conditional [control..., target] transfer sequence carrying every control item.", + ), + OutputParam( + "cond_no_control_static", + type_hint=dict, + kwargs_type="denoiser_input_fields", + description="Conditional [target] transfer sequence with the control items dropped.", + ), + OutputParam( + "uncond_full_static", + type_hint=dict, + kwargs_type="denoiser_input_fields", + description="Unconditional [control..., target] transfer sequence for text CFG.", + ), + OutputParam( + "num_noisy_vision_tokens", + type_hint=int, + description="Number of noisy target vision tokens denoised each step.", + ), + ] + + @torch.no_grad() + def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + device = components._execution_device + num_hints = len(block_state.control_latents) + + def _vision_pack(text_segment: dict, include_controls: bool) -> dict: + if include_controls: + vision_items = [*block_state.control_latents, block_state.latents] + condition_indexes = [None] * num_hints + [block_state.target_condition_indexes] + clean_flags = [True] * num_hints + [False] + else: + vision_items = [block_state.latents] + condition_indexes = [block_state.target_condition_indexes] + clean_flags = [False] + vision_segment = components._prepare_vision_segment( + input_vision_tokens=vision_items, + has_image_condition=False, + mrope_offset=text_segment["vision_start_temporal_offset"], + vision_fps=block_state.fps, + curr=text_segment["und_len"], + device=device, + condition_frame_indexes=condition_indexes, + clean_item_flags=clean_flags, + share_vision_temporal_positions=block_state.share_vision_temporal_positions, + ) + return { + **text_segment, + **vision_segment, + "position_ids": torch.cat([text_segment["text_mrope_ids"], vision_segment["vision_mrope_ids"]], dim=1), + "sequence_length": text_segment["und_len"] + vision_segment["num_vision_tokens"], + } + + block_state.cond_full_static = _vision_pack(block_state.cond_text_segment, include_controls=True) + block_state.cond_no_control_static = _vision_pack(block_state.cond_text_segment, include_controls=False) + block_state.uncond_full_static = _vision_pack(block_state.uncond_text_segment, include_controls=True) + block_state.num_noisy_vision_tokens = block_state.cond_full_static["num_noisy_vision_tokens"] + + self.set_block_state(state, block_state) + return components, state + + +class Cosmos3TransferSetTimestepsStep(ModularPipelineBlocks): + model_name = "cosmos3-omni" + + @property + def description(self) -> str: + return ( + "Resets the scheduler and computes timesteps for a single transfer chunk. UniPCMultistepScheduler keeps " + "per-step state on the instance, so it is reset per chunk (each autoregressive chunk is a full denoise)." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("scheduler", UniPCMultistepScheduler)] + + @property + def inputs(self) -> list[InputParam]: + return [InputParam.template("num_inference_steps", required=True)] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam("timesteps", type_hint=torch.Tensor, description="Scheduler timesteps for this chunk."), + OutputParam( + "num_warmup_steps", type_hint=int, description="Number of scheduler warmup steps for this chunk." + ), + ] + + @torch.no_grad() + def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + device = components._execution_device + components.scheduler.set_timesteps(block_state.num_inference_steps, device=device) + block_state.timesteps = components.scheduler.timesteps + block_state.num_warmup_steps = ( + len(block_state.timesteps) - block_state.num_inference_steps * components.scheduler.order + ) + self.set_block_state(state, block_state) + return components, state diff --git a/src/diffusers/modular_pipelines/cosmos/decoders.py b/src/diffusers/modular_pipelines/cosmos/decoders.py new file mode 100644 index 000000000000..7f3e4fa6393b --- /dev/null +++ b/src/diffusers/modular_pipelines/cosmos/decoders.py @@ -0,0 +1,256 @@ +import torch + +from ...configuration_utils import FrozenDict +from ...models.autoencoders.autoencoder_cosmos3_audio import Cosmos3AVAEAudioTokenizer +from ...models.autoencoders.autoencoder_kl_wan import AutoencoderKLWan +from ...utils import logging +from ...video_processor import VideoProcessor +from ..modular_pipeline import ModularPipelineBlocks, PipelineState +from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam +from .modular_pipeline import Cosmos3OmniModularPipeline + + +logger = logging.get_logger(__name__) + + +class Cosmos3VideoDecodeStep(ModularPipelineBlocks): + model_name = "cosmos3-omni" + + @property + def description(self) -> str: + return "Decodes denoised vision latents into video outputs." + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("vae", AutoencoderKLWan), + ComponentSpec( + "video_processor", + VideoProcessor, + config=FrozenDict({"vae_scale_factor": 16, "resample": "bilinear"}), + default_creation_method="from_config", + ), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("latents", required=True, description="Denoised vision latents to decode."), + InputParam.template("output_type", default="pil"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [OutputParam.template("videos")] + + @torch.no_grad() + def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + device = components._execution_device + + if block_state.output_type == "latent": + block_state.videos = block_state.latents + else: + in_dtype = block_state.latents.dtype + vae_dtype = components.vae.dtype + mean = components._vae_latents_mean.to(device=block_state.latents.device, dtype=vae_dtype) + inv_std = components._vae_latents_inv_std.to(device=block_state.latents.device, dtype=vae_dtype) + z_raw = block_state.latents.to(vae_dtype) / inv_std.view(1, -1, 1, 1, 1) + mean.view(1, -1, 1, 1, 1) + decoded = components.vae.decode(z_raw).sample.to(in_dtype) + block_state.videos = components.video_processor.postprocess_video( + decoded, output_type=block_state.output_type + )[0] + + if components.requires_safety_checker and block_state.output_type != "latent": + if getattr(components, "safety_checker", None) is None: + raise ValueError( + "Cosmos3 requires a safety checker by default. Call `pipe.enable_safety_checker()` to load it " + "(or pass your own), or opt out explicitly with `pipe.disable_safety_checker()`." + ) + block_state.videos = components._apply_video_safety_check( + block_state.videos, output_type=block_state.output_type, device=device + ) + + self.set_block_state(state, block_state) + return components, state + + +class Cosmos3SoundDecodeStep(ModularPipelineBlocks): + model_name = "cosmos3-omni" + + @property + def description(self) -> str: + return "Decodes sound latents into waveform output." + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("sound_tokenizer", Cosmos3AVAEAudioTokenizer)] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + name="sound_latents", + type_hint=torch.Tensor, + required=True, + description="Denoised sound latents to decode.", + ) + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam("sound", type_hint=torch.Tensor, description="Generated waveform."), + OutputParam("sampling_rate", type_hint=int, description="Sample rate of the generated waveform in Hz."), + ] + + @torch.no_grad() + def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + if components.sound_tokenizer is None: + raise ValueError("Sound decoding requires a sound-capable checkpoint with a sound_tokenizer.") + block_state.sound = components.decode_sound(block_state.sound_latents) + block_state.sampling_rate = int(components.sound_tokenizer.config.sampling_rate) + self.set_block_state(state, block_state) + return components, state + + +class Cosmos3TransferDecodeChunkStep(ModularPipelineBlocks): + model_name = "cosmos3-omni" + + @property + def description(self) -> str: + return ( + "Decodes one transfer chunk's latents to pixels (float32, clamped to [-1, 1]), records it as the " + "autoregressive seed for the next chunk, and appends it to output_chunks (dropping the overlap that " + "later chunks share with the previous chunk's conditioning frames)." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("vae", AutoencoderKLWan)] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + name="latents", type_hint=torch.Tensor, required=True, description="Denoised target latents for this chunk." + ), + InputParam(name="chunk_id", type_hint=int, default=0, description="Index of the current chunk."), + InputParam( + name="current_conditional_frames", + type_hint=int, + required=True, + description="Number of pixel frames this chunk reused from the previous chunk.", + ), + InputParam( + name="output_chunks", + type_hint=list[torch.Tensor], + required=True, + description="Decoded pixel chunks accumulated so far.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "previous_output", + type_hint=torch.Tensor, + description="Decoded pixels of this chunk, used to seed the next chunk.", + ), + OutputParam( + "output_chunks", + type_hint=list[torch.Tensor], + description="Decoded pixel chunks accumulated so far (with this chunk appended).", + ), + ] + + @torch.no_grad() + def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + latents = block_state.latents + vae_dtype = components.vae.dtype + mean = components._vae_latents_mean.to(device=latents.device, dtype=vae_dtype) + inv_std = components._vae_latents_inv_std.to(device=latents.device, dtype=vae_dtype) + z_raw = latents.to(vae_dtype) / inv_std.view(1, -1, 1, 1, 1) + mean.view(1, -1, 1, 1, 1) + output_video = components.vae.decode(z_raw).sample.to(torch.float32).clamp(-1, 1) + block_state.previous_output = output_video + chunk = ( + output_video if block_state.chunk_id == 0 else output_video[:, :, block_state.current_conditional_frames :] + ) + block_state.output_chunks = [*block_state.output_chunks, chunk] + self.set_block_state(state, block_state) + return components, state + + +class Cosmos3TransferStitchStep(ModularPipelineBlocks): + model_name = "cosmos3-omni" + + @property + def description(self) -> str: + return ( + "Concatenates the decoded transfer chunks along time, truncates to total_frames, and post-processes to " + "the requested output type. Transfer produces no audio, so sound / sampling_rate are None." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("vae", AutoencoderKLWan), + ComponentSpec( + "video_processor", + VideoProcessor, + config=FrozenDict({"vae_scale_factor": 16, "resample": "bilinear"}), + default_creation_method="from_config", + ), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + name="output_chunks", + type_hint=list[torch.Tensor], + required=True, + description="Decoded pixel chunks to stitch together.", + ), + InputParam( + name="total_frames", type_hint=int, required=True, description="Total number of output frames to keep." + ), + InputParam.template("output_type", default="pil"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam("videos", description="The generated transfer video."), + OutputParam("sound", description="Always None for transfer (no audio)."), + OutputParam("sampling_rate", description="Always None for transfer (no audio)."), + ] + + @torch.no_grad() + def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + device = components._execution_device + + decoded = torch.cat(block_state.output_chunks, dim=2)[:, :, : block_state.total_frames] + block_state.videos = components.video_processor.postprocess_video( + decoded, output_type=block_state.output_type + )[0] + + if components.requires_safety_checker and block_state.output_type != "latent": + if getattr(components, "safety_checker", None) is None: + raise ValueError( + "Cosmos3 requires a safety checker by default. Call `pipe.enable_safety_checker()` to load it " + "(or pass your own), or opt out explicitly with `pipe.disable_safety_checker()`." + ) + block_state.videos = components._apply_video_safety_check( + block_state.videos, output_type=block_state.output_type, device=device + ) + + block_state.sound = None + block_state.sampling_rate = None + self.set_block_state(state, block_state) + return components, state diff --git a/src/diffusers/modular_pipelines/cosmos/denoise.py b/src/diffusers/modular_pipelines/cosmos/denoise.py new file mode 100644 index 000000000000..06f21894f236 --- /dev/null +++ b/src/diffusers/modular_pipelines/cosmos/denoise.py @@ -0,0 +1,754 @@ +import inspect + +import torch + +from ...models.transformers.transformer_cosmos3 import Cosmos3OmniTransformer +from ...schedulers import UniPCMultistepScheduler +from ..modular_pipeline import ( + BlockState, + LoopSequentialPipelineBlocks, + ModularPipelineBlocks, + PipelineState, +) +from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam +from .modular_pipeline import Cosmos3OmniModularPipeline + + +class Cosmos3VisionLoopPrepareStep(ModularPipelineBlocks): + model_name = "cosmos3-omni" + + @property + def description(self) -> str: + return "Prepares vision tokens and timesteps for one denoising iteration." + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("transformer", Cosmos3OmniTransformer)] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("latents", required=True, description="Noisy vision latents to denoise."), + InputParam( + name="cond_vision_segment", type_hint=dict, required=True, description="Conditional vision segment." + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "vision_tokens", + type_hint=list[torch.Tensor], + description="Vision tokens for the transformer denoiser.", + ), + OutputParam("vision_timesteps", type_hint=torch.Tensor, description="Timesteps for the vision tokens."), + ] + + @torch.no_grad() + def __call__(self, components: Cosmos3OmniModularPipeline, block_state: BlockState, i: int, t: torch.Tensor): + device = components._execution_device + block_state.vision_tokens = [block_state.latents.to(device=device, dtype=components.transformer.dtype)] + block_state.vision_timesteps = torch.full( + (block_state.cond_vision_segment["num_noisy_vision_tokens"],), t.item(), device=device + ) + return components, block_state + + +class Cosmos3SoundLoopPrepareStep(ModularPipelineBlocks): + model_name = "cosmos3-omni" + + @property + def description(self) -> str: + return "Prepares sound tokens and timesteps for one denoising iteration." + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("transformer", Cosmos3OmniTransformer)] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + name="sound_latents", + type_hint=torch.Tensor, + required=True, + description="Noisy sound latents to denoise.", + ), + InputParam( + name="cond_sound_segment", type_hint=dict, required=True, description="Conditional sound segment." + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "sound_tokens", type_hint=list[torch.Tensor], description="Sound tokens for the transformer denoiser." + ), + OutputParam("sound_timesteps", type_hint=torch.Tensor, description="Timesteps for the sound tokens."), + ] + + @torch.no_grad() + def __call__(self, components: Cosmos3OmniModularPipeline, block_state: BlockState, i: int, t: torch.Tensor): + device = components._execution_device + block_state.sound_tokens = [block_state.sound_latents.to(device=device, dtype=components.transformer.dtype)] + block_state.sound_timesteps = torch.full( + (block_state.cond_sound_segment["sound_len"],), t.item(), device=device + ) + return components, block_state + + +class Cosmos3ActionLoopPrepareStep(ModularPipelineBlocks): + model_name = "cosmos3-omni" + + @property + def description(self) -> str: + return "Prepares action tokens and timesteps for one denoising iteration." + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("transformer", Cosmos3OmniTransformer)] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + name="action_latents", + type_hint=torch.Tensor, + required=True, + description="Noisy action latents to denoise.", + ), + InputParam( + name="cond_action_segment", type_hint=dict, required=True, description="Conditional action segment." + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "action_tokens", + type_hint=list[torch.Tensor], + description="Action tokens for the transformer denoiser.", + ), + OutputParam("action_timesteps", type_hint=torch.Tensor, description="Timesteps for the action tokens."), + ] + + @torch.no_grad() + def __call__(self, components: Cosmos3OmniModularPipeline, block_state: BlockState, i: int, t: torch.Tensor): + device = components._execution_device + block_state.action_tokens = [block_state.action_latents.to(device=device, dtype=components.transformer.dtype)] + block_state.action_timesteps = torch.full( + (block_state.cond_action_segment["num_noisy_action_tokens"],), t.item(), device=device + ) + return components, block_state + + +class Cosmos3LoopDenoiser(ModularPipelineBlocks): + model_name = "cosmos3-omni" + + @property + def description(self) -> str: + return "Predicts available Cosmos3 modality velocities for one denoising iteration." + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("transformer", Cosmos3OmniTransformer)] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("denoiser_input_fields"), + InputParam( + name="guidance_scale", + type_hint=float, + default=6.0, + description="Scale for classifier-free guidance.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "velocity_vision", type_hint=torch.Tensor, description="Predicted velocity for vision latents." + ), + OutputParam("velocity_sound", type_hint=torch.Tensor, description="Predicted velocity for sound latents."), + OutputParam( + "velocity_action", type_hint=torch.Tensor, description="Predicted velocity for action latents." + ), + ] + + @torch.no_grad() + def __call__(self, components: Cosmos3OmniModularPipeline, block_state: BlockState, i: int, t: torch.Tensor): + denoiser_input_fields = block_state.denoiser_input_fields + loop_input_fields = block_state.as_dict() + has_sound = "sound_tokens" in loop_input_fields + has_action = "action_tokens" in loop_input_fields + do_cfg = block_state.guidance_scale != 1.0 + transformer_args = set(inspect.signature(components.transformer.forward).parameters) + + prediction_passes = ["cond"] + if do_cfg: + prediction_passes.append("uncond") + + velocities = {} + for pass_name in prediction_passes: + transformer_kwargs = {} + for field_name, field_value in denoiser_input_fields.items(): + if field_name.startswith(f"{pass_name}_"): + transformer_field_name = field_name.removeprefix(f"{pass_name}_") + if transformer_field_name.endswith("_segment"): + transformer_kwargs.update(field_value) + else: + transformer_kwargs[transformer_field_name] = field_value + elif field_name in transformer_args: + transformer_kwargs[field_name] = field_value + transformer_kwargs.update( + { + field_name: field_value + for field_name, field_value in loop_input_fields.items() + if field_name in transformer_args + } + ) + transformer_kwargs = { + name: value for name, value in transformer_kwargs.items() if name in transformer_args + } + preds_vision, preds_sound, preds_action = components.transformer(**transformer_kwargs) + velocities[pass_name] = components._mask_velocity_predictions( + preds_vision, + preds_sound, + vision_condition_mask=[loop_input_fields["vision_condition_mask"]], + sound_condition_mask=[loop_input_fields["sound_condition_mask"]] if has_sound else None, + preds_action=preds_action, + action_condition_mask=[loop_input_fields["action_condition_mask"]] if has_action else None, + raw_action_dim=loop_input_fields.get("raw_action_dim_resolved"), + ) + + cond_velocity_vision, cond_velocity_sound, cond_velocity_action = velocities["cond"] + if do_cfg: + uncond_velocity_vision, uncond_velocity_sound, uncond_velocity_action = velocities["uncond"] + block_state.velocity_vision = uncond_velocity_vision + block_state.guidance_scale * ( + cond_velocity_vision - uncond_velocity_vision + ) + block_state.velocity_sound = ( + uncond_velocity_sound + block_state.guidance_scale * (cond_velocity_sound - uncond_velocity_sound) + if has_sound + else None + ) + block_state.velocity_action = ( + uncond_velocity_action + block_state.guidance_scale * (cond_velocity_action - uncond_velocity_action) + if has_action + else None + ) + else: + block_state.velocity_vision = cond_velocity_vision + block_state.velocity_sound = cond_velocity_sound if has_sound else None + block_state.velocity_action = cond_velocity_action if has_action else None + + return components, block_state + + +class Cosmos3VisionLoopSchedulerStep(ModularPipelineBlocks): + model_name = "cosmos3-omni" + + @property + def description(self) -> str: + return "Updates vision latents after one denoising iteration." + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("scheduler", UniPCMultistepScheduler)] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("latents", required=True, description="Noisy vision latents to update."), + InputParam( + name="velocity_vision", type_hint=torch.Tensor, required=True, description="Predicted vision velocity." + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [OutputParam.template("latents")] + + @torch.no_grad() + def __call__(self, components: Cosmos3OmniModularPipeline, block_state: BlockState, i: int, t: torch.Tensor): + block_state.latents = components.scheduler.step( + block_state.velocity_vision.unsqueeze(0), t, block_state.latents.unsqueeze(0), return_dict=False + )[0].squeeze(0) + return components, block_state + + +class Cosmos3SoundLoopSchedulerStep(ModularPipelineBlocks): + model_name = "cosmos3-omni" + + @property + def description(self) -> str: + return "Updates sound latents after one denoising iteration." + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + name="sound_latents", + type_hint=torch.Tensor, + required=True, + description="Noisy sound latents to update.", + ), + InputParam( + name="sound_scheduler", + type_hint=UniPCMultistepScheduler, + required=True, + description="Scheduler used to update sound latents.", + ), + InputParam( + name="velocity_sound", type_hint=torch.Tensor, required=True, description="Predicted sound velocity." + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [OutputParam("sound_latents", type_hint=torch.Tensor, description="Updated sound latents.")] + + @torch.no_grad() + def __call__(self, components: Cosmos3OmniModularPipeline, block_state: BlockState, i: int, t: torch.Tensor): + block_state.sound_latents = block_state.sound_scheduler.step( + block_state.velocity_sound.unsqueeze(0), t, block_state.sound_latents.unsqueeze(0), return_dict=False + )[0].squeeze(0) + return components, block_state + + +class Cosmos3ActionLoopSchedulerStep(ModularPipelineBlocks): + model_name = "cosmos3-omni" + + @property + def description(self) -> str: + return "Updates action latents after one denoising iteration." + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + name="action_latents", + type_hint=torch.Tensor, + required=True, + description="Noisy action latents to update.", + ), + InputParam( + name="action_scheduler", + type_hint=UniPCMultistepScheduler, + required=True, + description="Scheduler used to update action latents.", + ), + InputParam( + name="velocity_action", type_hint=torch.Tensor, required=True, description="Predicted action velocity." + ), + InputParam( + name="action_condition_mask", + type_hint=torch.Tensor, + required=True, + description="Mask marking conditioned action latent frames.", + ), + InputParam( + name="raw_action_dim_resolved", + type_hint=int, + default=None, + description="Unpadded action-vector dimension.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [OutputParam("action_latents", type_hint=torch.Tensor, description="Updated action latents.")] + + @torch.no_grad() + def __call__(self, components: Cosmos3OmniModularPipeline, block_state: BlockState, i: int, t: torch.Tensor): + has_noisy_action = block_state.action_condition_mask.sum() < block_state.action_condition_mask.numel() + if has_noisy_action: + block_state.action_latents = block_state.action_scheduler.step( + block_state.velocity_action.unsqueeze(0), t, block_state.action_latents.unsqueeze(0), return_dict=False + )[0].squeeze(0) + if block_state.raw_action_dim_resolved is not None: + block_state.action_latents[:, block_state.raw_action_dim_resolved :] = 0 + return components, block_state + + +class Cosmos3DenoiseLoopWrapper(LoopSequentialPipelineBlocks): + model_name = "cosmos3-omni" + + @property + def description(self) -> str: + return "Iteratively denoises Cosmos3 latents over scheduler timesteps." + + @property + def loop_expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("scheduler", UniPCMultistepScheduler), + ComponentSpec("transformer", Cosmos3OmniTransformer), + ] + + @property + def loop_inputs(self) -> list[InputParam]: + return [ + InputParam.template("timesteps", required=True), + InputParam.template("num_inference_steps", required=True), + InputParam( + name="num_warmup_steps", type_hint=int, required=True, description="Number of scheduler warmup steps." + ), + ] + + @torch.no_grad() + def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + with self.progress_bar(total=block_state.num_inference_steps) as progress_bar: + for i, t in enumerate(block_state.timesteps): + components, block_state = self.loop_step(components, block_state, i=i, t=t) + if i == len(block_state.timesteps) - 1 or ( + (i + 1) > block_state.num_warmup_steps and (i + 1) % components.scheduler.order == 0 + ): + progress_bar.update() + self.set_block_state(state, block_state) + return components, state + + +class Cosmos3VisionDenoiseStep(Cosmos3DenoiseLoopWrapper): + block_classes = [ + Cosmos3VisionLoopPrepareStep, + Cosmos3LoopDenoiser, + Cosmos3VisionLoopSchedulerStep, + ] + block_names = ["prepare_vision", "denoiser", "update_vision"] + + @property + def description(self) -> str: + return "Runs the vision-only Cosmos3 denoising loop." + + +class Cosmos3VisionSoundDenoiseStep(Cosmos3DenoiseLoopWrapper): + block_classes = [ + Cosmos3VisionLoopPrepareStep, + Cosmos3SoundLoopPrepareStep, + Cosmos3LoopDenoiser, + Cosmos3VisionLoopSchedulerStep, + Cosmos3SoundLoopSchedulerStep, + ] + block_names = ["prepare_vision", "prepare_sound", "denoiser", "update_vision", "update_sound"] + + @property + def description(self) -> str: + return "Runs the vision-and-sound Cosmos3 denoising loop." + + +class Cosmos3VisionActionDenoiseStep(Cosmos3DenoiseLoopWrapper): + block_classes = [ + Cosmos3VisionLoopPrepareStep, + Cosmos3ActionLoopPrepareStep, + Cosmos3LoopDenoiser, + Cosmos3VisionLoopSchedulerStep, + Cosmos3ActionLoopSchedulerStep, + ] + block_names = ["prepare_vision", "prepare_action", "denoiser", "update_vision", "update_action"] + + @property + def description(self) -> str: + return "Runs the vision-and-action Cosmos3 denoising loop." + + +class Cosmos3VisionSoundActionDenoiseStep(Cosmos3DenoiseLoopWrapper): + block_classes = [ + Cosmos3VisionLoopPrepareStep, + Cosmos3SoundLoopPrepareStep, + Cosmos3ActionLoopPrepareStep, + Cosmos3LoopDenoiser, + Cosmos3VisionLoopSchedulerStep, + Cosmos3SoundLoopSchedulerStep, + Cosmos3ActionLoopSchedulerStep, + ] + block_names = [ + "prepare_vision", + "prepare_sound", + "prepare_action", + "denoiser", + "update_vision", + "update_sound", + "update_action", + ] + + @property + def description(self) -> str: + return "Runs the vision, sound, and action Cosmos3 denoising loop." + + +class Cosmos3TransferLoopPrepareStep(ModularPipelineBlocks): + model_name = "cosmos3-omni" + + @property + def description(self) -> str: + return "Prepares the full [control..., target] and target-only vision token lists plus timesteps for one transfer iteration." + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("transformer", Cosmos3OmniTransformer)] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + name="control_latents", + type_hint=list[torch.Tensor], + required=True, + description="Clean control latents for this chunk, one per hint in canonical order.", + ), + InputParam( + name="latents", type_hint=torch.Tensor, required=True, description="Noisy target latents to denoise." + ), + InputParam( + name="num_noisy_vision_tokens", + type_hint=int, + required=True, + description="Number of noisy target vision tokens denoised each step.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "vision_tokens_full", + type_hint=list[torch.Tensor], + description="Token list for the [control..., target] forward passes.", + ), + OutputParam( + "vision_tokens_target", + type_hint=list[torch.Tensor], + description="Token list for the target-only (no-control) forward pass.", + ), + OutputParam( + "vision_timesteps", type_hint=torch.Tensor, description="Timesteps for the noisy target tokens." + ), + ] + + @torch.no_grad() + def __call__(self, components: Cosmos3OmniModularPipeline, block_state: BlockState, i: int, t: torch.Tensor): + device = components._execution_device + dtype = components.transformer.dtype + block_state.vision_tokens_full = [c.to(device=device, dtype=dtype) for c in block_state.control_latents] + [ + block_state.latents.to(device=device, dtype=dtype) + ] + block_state.vision_tokens_target = [block_state.latents.to(device=device, dtype=dtype)] + block_state.vision_timesteps = torch.full((block_state.num_noisy_vision_tokens,), t.item(), device=device) + return components, block_state + + +class Cosmos3TransferLoopDenoiser(ModularPipelineBlocks): + # Dedicated (not Cosmos3LoopDenoiser): transfer runs up to 3 passes over different token sequences with nested + # control/text CFG and interval gating, which the generic cond/uncond denoiser cannot express. + model_name = "cosmos3-omni" + + @property + def description(self) -> str: + return ( + "Predicts the transfer velocity with nested control/text CFG over [control..., target]. Each branch is " + "gated by its guidance interval, and the result is masked so conditioned frames get zero velocity." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("transformer", Cosmos3OmniTransformer)] + + @property + def inputs(self) -> list[InputParam]: + return [ + # The three pre-packed CFG sequence variants (cond_full / cond_no_control / uncond_full) flow in as + # denoiser_input_fields, gathered generically like the other Cosmos3 denoisers. + InputParam.template("denoiser_input_fields"), + InputParam( + name="vision_tokens_full", + type_hint=list[torch.Tensor], + required=True, + description="Token list for the [control..., target] forward passes.", + ), + InputParam( + name="vision_tokens_target", + type_hint=list[torch.Tensor], + required=True, + description="Token list for the target-only (no-control) forward pass.", + ), + InputParam( + name="vision_timesteps", + type_hint=torch.Tensor, + required=True, + description="Timesteps for the noisy target tokens.", + ), + InputParam( + name="velocity_mask", + type_hint=torch.Tensor, + required=True, + description="Mask that zeroes the velocity on conditioned (clean) latent frames.", + ), + InputParam( + name="guidance_scale", type_hint=float, default=6.0, description="Scale for text classifier-free guidance." + ), + InputParam( + name="control_guidance", + type_hint=float, + default=1.0, + description="Scale for the control (structural) guidance axis.", + ), + InputParam( + name="guidance_interval", + type_hint=tuple, + default=None, + description="Timestep interval [lo, hi] over which text guidance is active (None = always).", + ), + InputParam( + name="control_guidance_interval", + type_hint=tuple, + default=None, + description="Timestep interval [lo, hi] over which control guidance is active (None = always).", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [OutputParam("velocity", type_hint=torch.Tensor, description="Predicted (masked) transfer velocity.")] + + @staticmethod + def _forward(components, static, vision_tokens, vision_timesteps): + preds_vision, _, _ = components.transformer( + input_ids=static["input_ids"], + text_indexes=static["text_indexes"], + position_ids=static["position_ids"], + und_len=static["und_len"], + sequence_length=static["sequence_length"], + vision_tokens=vision_tokens, + vision_token_shapes=static["vision_token_shapes"], + vision_sequence_indexes=static["vision_sequence_indexes"], + vision_mse_loss_indexes=static["vision_mse_loss_indexes"], + vision_timesteps=vision_timesteps, + vision_noisy_frame_indexes=static["vision_noisy_frame_indexes"], + ) + return preds_vision[-1] + + @torch.no_grad() + def __call__(self, components: Cosmos3OmniModularPipeline, block_state: BlockState, i: int, t: torch.Tensor): + # active-at: a None interval is always active; otherwise the timestep must fall within [lo, hi]. + guidance_interval = block_state.guidance_interval + guidance_active = guidance_interval is None or ( + float(guidance_interval[0]) <= float(t.item()) <= float(guidance_interval[1]) + ) + control_interval = block_state.control_guidance_interval + control_active = control_interval is None or ( + float(control_interval[0]) <= float(t.item()) <= float(control_interval[1]) + ) + step_guidance = block_state.guidance_scale if guidance_active else 1.0 + step_control = block_state.control_guidance if control_active else 1.0 + needs_text_cfg = step_guidance > 1.0 + needs_control_cfg = step_control != 1.0 + + denoiser_input_fields = block_state.denoiser_input_fields + cond_full_static = denoiser_input_fields["cond_full_static"] + cond_no_control_static = denoiser_input_fields["cond_no_control_static"] + uncond_full_static = denoiser_input_fields["uncond_full_static"] + + cond_full = self._forward( + components, cond_full_static, block_state.vision_tokens_full, block_state.vision_timesteps + ) + + cond_no_control = None + if needs_control_cfg: + cond_no_control = self._forward( + components, + cond_no_control_static, + block_state.vision_tokens_target, + block_state.vision_timesteps, + ) + + uncond_full = None + if needs_text_cfg: + uncond_full = self._forward( + components, + uncond_full_static, + block_state.vision_tokens_full, + block_state.vision_timesteps, + ) + + if needs_control_cfg and needs_text_cfg: + control_cond = cond_no_control + step_control * (cond_full - cond_no_control) + velocity = uncond_full + step_guidance * (control_cond - uncond_full) + elif needs_control_cfg: + velocity = cond_no_control + step_control * (cond_full - cond_no_control) + elif needs_text_cfg: + velocity = uncond_full + step_guidance * (cond_full - uncond_full) + else: + velocity = cond_full + + block_state.velocity = velocity * block_state.velocity_mask + return components, block_state + + +class Cosmos3TransferLoopSchedulerStep(ModularPipelineBlocks): + model_name = "cosmos3-omni" + + @property + def description(self) -> str: + return "Steps the scheduler and re-pins the conditioned frames exactly for one transfer iteration." + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("scheduler", UniPCMultistepScheduler)] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + name="latents", type_hint=torch.Tensor, required=True, description="Noisy target latents to update." + ), + InputParam( + name="velocity", type_hint=torch.Tensor, required=True, description="Predicted (masked) transfer velocity." + ), + InputParam( + name="velocity_mask", + type_hint=torch.Tensor, + required=True, + description="Mask that zeroes the velocity on conditioned (clean) latent frames.", + ), + InputParam( + name="condition_latents", + type_hint=torch.Tensor, + required=True, + description="Clean target latents on the conditioned frames (the autoregressive seed).", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [OutputParam("latents", type_hint=torch.Tensor, description="Updated target latents for this chunk.")] + + @torch.no_grad() + def __call__(self, components: Cosmos3OmniModularPipeline, block_state: BlockState, i: int, t: torch.Tensor): + block_state.latents = components.scheduler.step( + block_state.velocity.unsqueeze(0), t, block_state.latents.unsqueeze(0), return_dict=False + )[0].squeeze(0) + # Re-pin conditioned frames exactly (the autoregressive seed), guarding multistep drift. + block_state.latents = ( + block_state.velocity_mask * block_state.latents + + (1.0 - block_state.velocity_mask) * block_state.condition_latents + ) + return components, block_state + + +class Cosmos3TransferDenoiseStep(Cosmos3DenoiseLoopWrapper): + block_classes = [ + Cosmos3TransferLoopPrepareStep, + Cosmos3TransferLoopDenoiser, + Cosmos3TransferLoopSchedulerStep, + ] + block_names = ["prepare_transfer", "denoiser", "update_transfer"] + + @property + def description(self) -> str: + return "Runs the per-chunk transfer denoising loop over scheduler timesteps." diff --git a/src/diffusers/modular_pipelines/cosmos/encoders.py b/src/diffusers/modular_pipelines/cosmos/encoders.py new file mode 100644 index 000000000000..9b3f0e5c8dcf --- /dev/null +++ b/src/diffusers/modular_pipelines/cosmos/encoders.py @@ -0,0 +1,913 @@ +import torch +from transformers import AutoTokenizer + +from ...configuration_utils import FrozenDict +from ...models.autoencoders.autoencoder_kl_wan import AutoencoderKLWan +from ...pipelines.cosmos.pipeline_cosmos3_omni import ( + _ACTION_RESOLUTION_BINS, + CosmosActionCondition, +) +from ...utils import logging +from ...video_processor import VideoProcessor +from ..modular_pipeline import ModularPipelineBlocks, PipelineState +from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam +from .modular_pipeline import Cosmos3OmniModularPipeline + + +logger = logging.get_logger(__name__) + + +class Cosmos3TextEncoderStep(ModularPipelineBlocks): + model_name = "cosmos3-omni" + + @property + def description(self) -> str: + return "Prepares non-action prompt token IDs for downstream text-segment packing." + + @staticmethod + def _check_inputs(block_state) -> None: + prompt = block_state.prompt + negative_prompt = block_state.negative_prompt + + if not isinstance(prompt, str): + raise ValueError( + f"`prompt` must be a str; batched prompts are not supported, got {type(prompt).__name__}." + ) + if negative_prompt is not None and not isinstance(negative_prompt, str): + raise ValueError( + "`negative_prompt` must be a str or None; batched prompts are not supported, " + f"got {type(negative_prompt).__name__}." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("text_tokenizer", AutoTokenizer), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("prompt", description="The text prompt that guides Cosmos3 generation."), + InputParam.template( + "negative_prompt", description="The negative text prompt used for classifier-free guidance." + ), + InputParam(name="num_frames", type_hint=int, default=None, description="Number of frames to generate."), + InputParam( + name="height", + type_hint=int, + default=None, + description="Height of the generated video or image in pixels.", + ), + InputParam( + name="width", + type_hint=int, + default=None, + description="Width of the generated video or image in pixels.", + ), + InputParam(name="fps", type_hint=float, default=24.0, description="Frame rate of the generated video."), + InputParam( + name="use_system_prompt", + type_hint=bool, + default=True, + description="Whether to prepend the Cosmos3 system prompt.", + ), + InputParam( + name="add_resolution_template", + type_hint=bool, + default=True, + description="Whether to add resolution metadata to the prompt.", + ), + InputParam( + name="add_duration_template", + type_hint=bool, + default=True, + description="Whether to add duration metadata to the prompt.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam("num_frames", type_hint=int, description="Number of frames to generate."), + OutputParam("height", type_hint=int, description="Height of the generated video or image in pixels."), + OutputParam("width", type_hint=int, description="Width of the generated video or image in pixels."), + OutputParam("cond_input_ids", type_hint=torch.Tensor, description="Token IDs for the conditional prompt."), + OutputParam( + "uncond_input_ids", type_hint=torch.Tensor, description="Token IDs for the unconditional prompt." + ), + ] + + @torch.no_grad() + def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + if block_state.num_frames is None: + block_state.num_frames = 189 + if block_state.height is None: + block_state.height = 720 + if block_state.width is None: + block_state.width = 1280 + + self._check_inputs(block_state) + if components.requires_safety_checker: + if getattr(components, "safety_checker", None) is None: + raise ValueError( + "Cosmos3 requires a safety checker by default. Call `pipe.enable_safety_checker()` to load it " + "(or pass your own), or opt out explicitly with `pipe.disable_safety_checker()`." + ) + device = components._execution_device + components.safety_checker.to(device) + try: + if not components.safety_checker.check_text_safety(block_state.prompt): + raise ValueError( + f"Cosmos Guardrail detected unsafe text in the prompt: {block_state.prompt}. " + "Please ensure that the prompt abides by the NVIDIA Open Model License Agreement." + ) + finally: + components.safety_checker.to("cpu") + + block_state.cond_input_ids, block_state.uncond_input_ids = components.tokenize_prompt( + block_state.prompt, + block_state.negative_prompt, + num_frames=block_state.num_frames, + height=block_state.height, + width=block_state.width, + fps=block_state.fps, + use_system_prompt=block_state.use_system_prompt, + add_resolution_template=block_state.add_resolution_template, + add_duration_template=block_state.add_duration_template, + action_mode=None, + action_view_point=None, + ) + + self.set_block_state(state, block_state) + return components, state + + +class Cosmos3TransferTextStep(ModularPipelineBlocks): + model_name = "cosmos3-omni" + + @property + def description(self) -> str: + return ( + "Tokenizes the transfer prompt in transfer mode using the per-chunk frame count. Transfer prompts are " + "pre-upsampled JSON captions passed through verbatim (the metadata templates are skipped)." + ) + + @staticmethod + def _check_inputs(block_state) -> None: + prompt = block_state.prompt + negative_prompt = block_state.negative_prompt + + if not isinstance(prompt, (str, list)) or ( + isinstance(prompt, list) and not all(isinstance(p, str) for p in prompt) + ): + raise ValueError(f"`prompt` must be a str or list of str, got {type(prompt).__name__}.") + if negative_prompt is not None and not isinstance(negative_prompt, (str, list)): + raise ValueError( + f"`negative_prompt` must be a str, list of str, or None, got {type(negative_prompt).__name__}." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("text_tokenizer", AutoTokenizer), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + name="prompt", + type_hint=str, + required=True, + description="The text prompt that guides Cosmos3 generation.", + ), + InputParam( + name="negative_prompt", + type_hint=str, + default=None, + description="The negative text prompt used for classifier-free guidance.", + ), + InputParam( + name="chunk_frames", type_hint=int, required=True, description="Number of pixel frames in this chunk." + ), + InputParam( + name="height", + type_hint=int, + default=None, + description="Height of the generated video in pixels.", + ), + InputParam( + name="width", type_hint=int, default=None, description="Width of the generated video in pixels." + ), + InputParam(name="fps", type_hint=float, default=24.0, description="Frame rate of the generated video."), + InputParam( + name="use_system_prompt", + type_hint=bool, + default=True, + description="Whether to prepend the Cosmos3 system prompt.", + ), + InputParam( + name="add_resolution_template", + type_hint=bool, + default=True, + description="Whether to add resolution metadata to the prompt.", + ), + InputParam( + name="add_duration_template", + type_hint=bool, + default=True, + description="Whether to add duration metadata to the prompt.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam("cond_input_ids", type_hint=torch.Tensor, description="Token IDs for the conditional prompt."), + OutputParam( + "uncond_input_ids", type_hint=torch.Tensor, description="Token IDs for the unconditional prompt." + ), + ] + + @torch.no_grad() + def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + self._check_inputs(block_state) + + if isinstance(block_state.prompt, list): + block_state.prompt = block_state.prompt[0] + if isinstance(block_state.negative_prompt, list): + block_state.negative_prompt = block_state.negative_prompt[0] + + if components.requires_safety_checker: + if getattr(components, "safety_checker", None) is None: + raise ValueError( + "Cosmos3 requires a safety checker by default. Call `pipe.enable_safety_checker()` to load it " + "(or pass your own), or opt out explicitly with `pipe.disable_safety_checker()`." + ) + device = components._execution_device + components.safety_checker.to(device) + try: + if not components.safety_checker.check_text_safety(block_state.prompt): + raise ValueError( + f"Cosmos Guardrail detected unsafe text in the prompt: {block_state.prompt}. " + "Please ensure that the prompt abides by the NVIDIA Open Model License Agreement." + ) + finally: + components.safety_checker.to("cpu") + + block_state.cond_input_ids, block_state.uncond_input_ids = components.tokenize_prompt( + block_state.prompt, + block_state.negative_prompt, + num_frames=block_state.chunk_frames, + height=block_state.height, + width=block_state.width, + fps=block_state.fps, + use_system_prompt=block_state.use_system_prompt, + add_resolution_template=block_state.add_resolution_template, + add_duration_template=block_state.add_duration_template, + action_mode=None, + action_view_point=None, + transfer_mode=True, + ) + + self.set_block_state(state, block_state) + return components, state + + +class Cosmos3ActionTextStep(ModularPipelineBlocks): + model_name = "cosmos3-omni" + + @property + def description(self) -> str: + return "Prepares action prompt token IDs from prompt + action metadata." + + @staticmethod + def _check_inputs(block_state) -> None: + prompt = block_state.prompt + negative_prompt = block_state.negative_prompt + action = block_state.action + num_frames = block_state.num_frames + height = block_state.height + width = block_state.width + if not isinstance(prompt, str): + raise ValueError( + f"`prompt` must be a str; batched prompts are not supported, got {type(prompt).__name__}." + ) + if negative_prompt is not None and not isinstance(negative_prompt, str): + raise ValueError( + "`negative_prompt` must be a str or None; batched prompts are not supported, " + f"got {type(negative_prompt).__name__}." + ) + if action is None: + raise ValueError("`action` is required for Cosmos3ActionTextStep.") + if action.image is None and action.video is None: + raise ValueError("`action.image` or `action.video` must be provided for action-conditioned generation.") + if num_frames is not None: + raise ValueError("`num_frames` has to be None if action is not None.") + if height is not None or width is not None: + raise ValueError("`height` and `width` have to be None if action is not None.") + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("text_tokenizer", AutoTokenizer), + ComponentSpec( + "video_processor", + VideoProcessor, + config=FrozenDict({"vae_scale_factor": 16, "resample": "bilinear"}), + default_creation_method="from_config", + ), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("prompt", description="The text prompt that guides Cosmos3 generation."), + InputParam.template( + "negative_prompt", description="The negative text prompt used for classifier-free guidance." + ), + InputParam( + name="action", + type_hint=CosmosActionCondition, + required=True, + description="Action-conditioning metadata and its reference visual input.", + ), + InputParam(name="num_frames", type_hint=int, default=None, description="Number of frames to generate."), + InputParam( + name="height", + type_hint=int, + default=None, + description="Height of the generated video or image in pixels.", + ), + InputParam( + name="width", + type_hint=int, + default=None, + description="Width of the generated video or image in pixels.", + ), + InputParam(name="fps", type_hint=float, default=24.0, description="Frame rate of the generated video."), + InputParam( + name="use_system_prompt", + type_hint=bool, + default=True, + description="Whether to prepend the Cosmos3 system prompt.", + ), + InputParam( + name="add_resolution_template", + type_hint=bool, + default=True, + description="Whether to add resolution metadata to the prompt.", + ), + InputParam( + name="add_duration_template", + type_hint=bool, + default=True, + description="Whether to add duration metadata to the prompt.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam("action_mode", type_hint=str, description="Requested action-generation mode."), + OutputParam("num_frames", type_hint=int, description="Number of frames to generate."), + OutputParam("height", type_hint=int, description="Height of the generated video or image in pixels."), + OutputParam("width", type_hint=int, description="Width of the generated video or image in pixels."), + OutputParam("cond_input_ids", type_hint=torch.Tensor, description="Token IDs for the conditional prompt."), + OutputParam( + "uncond_input_ids", type_hint=torch.Tensor, description="Token IDs for the unconditional prompt." + ), + ] + + @torch.no_grad() + def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + self._check_inputs(block_state) + + action = block_state.action + block_state.action_mode = action.mode + block_state.num_frames = action.chunk_size + 1 + conditioning_clip = [action.image] if action.image is not None else action.video + probe = components.video_processor.preprocess_video(conditioning_clip) + source_h, source_w = int(probe.shape[-2]), int(probe.shape[-1]) + resolution_key = str(action.resolution_tier) + block_state.height, block_state.width = VideoProcessor.classify_height_width_bin( + source_h, source_w, ratios=_ACTION_RESOLUTION_BINS[resolution_key] + ) + + if components.requires_safety_checker: + if getattr(components, "safety_checker", None) is None: + raise ValueError( + "Cosmos3 requires a safety checker by default. Call `pipe.enable_safety_checker()` to load it " + "(or pass your own), or opt out explicitly with `pipe.disable_safety_checker()`." + ) + device = components._execution_device + components.safety_checker.to(device) + try: + if not components.safety_checker.check_text_safety(block_state.prompt): + raise ValueError( + f"Cosmos Guardrail detected unsafe text in the prompt: {block_state.prompt}. " + "Please ensure that the prompt abides by the NVIDIA Open Model License Agreement." + ) + finally: + components.safety_checker.to("cpu") + + block_state.cond_input_ids, block_state.uncond_input_ids = components.tokenize_prompt( + block_state.prompt, + block_state.negative_prompt, + num_frames=block_state.num_frames, + height=block_state.height, + width=block_state.width, + fps=block_state.fps, + use_system_prompt=block_state.use_system_prompt, + add_resolution_template=block_state.add_resolution_template, + add_duration_template=block_state.add_duration_template, + action_mode=block_state.action_mode, + action_view_point=action.view_point, + ) + + self.set_block_state(state, block_state) + return components, state + + +class Cosmos3ImageVaeEncoderStep(ModularPipelineBlocks): + model_name = "cosmos3-omni" + + @property + def description(self) -> str: + return "Encodes non-action image-to-video conditioning into Cosmos3 vision latents." + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("vae", AutoencoderKLWan), + ComponentSpec( + "video_processor", + VideoProcessor, + config=FrozenDict({"vae_scale_factor": 16, "resample": "bilinear"}), + default_creation_method="from_config", + ), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam(name="image", default=None, description="Reference image for image-to-video conditioning."), + InputParam(name="num_frames", type_hint=int, required=True, description="Number of frames to generate."), + InputParam( + name="height", type_hint=int, required=True, description="Height of the generated video in pixels." + ), + InputParam( + name="width", type_hint=int, required=True, description="Width of the generated video in pixels." + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "x0_tokens_vision", + type_hint=torch.Tensor, + description="Vision latents encoded from the conditioning image or video.", + ), + OutputParam( + "vision_condition_frames", + type_hint=list[int], + description="Latent-frame indexes fixed by visual conditioning.", + ), + ] + + @torch.no_grad() + def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + device = components._execution_device + dtype = components.transformer.dtype + + if block_state.image is None: + raise ValueError("`Cosmos3ImageVaeEncoderStep` requires an `image` input.") + if block_state.num_frames == 1: + raise ValueError( + "`image` conditioning requires `num_frames` > 1; image-to-image generation is not supported." + ) + if block_state.num_frames < 1: + raise ValueError(f"`num_frames` must be >= 1, got {block_state.num_frames}.") + + sf = int(components.vae.config.scale_factor_spatial) + if block_state.height % sf != 0 or block_state.width % sf != 0: + raise ValueError( + f"`height` and `width` must be multiples of {sf}, got ({block_state.height}, {block_state.width})." + ) + + conditioning_frame_2d = components.video_processor.preprocess( + block_state.image, height=block_state.height, width=block_state.width + ).to(device=device, dtype=dtype) + + vision_tensor = torch.zeros( + 1, + 3, + block_state.num_frames, + block_state.height, + block_state.width, + dtype=dtype, + device=device, + ) + vision_tensor[:, :, 0] = conditioning_frame_2d + vision_tensor[:, :, 1:] = conditioning_frame_2d.unsqueeze(2).expand(-1, -1, block_state.num_frames - 1, -1, -1) + + block_state.x0_tokens_vision = components._encode_video(vision_tensor).contiguous().float() + block_state.vision_condition_frames = [0] + + self.set_block_state(state, block_state) + return components, state + + +class Cosmos3VideoVaeEncoderStep(ModularPipelineBlocks): + model_name = "cosmos3-omni" + + @property + def description(self) -> str: + return "Encodes non-action video conditioning into Cosmos3 vision latents." + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("vae", AutoencoderKLWan), + ComponentSpec( + "video_processor", + VideoProcessor, + config=FrozenDict({"vae_scale_factor": 16, "resample": "bilinear"}), + default_creation_method="from_config", + ), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam(name="video", default=None, description="Reference video for video-to-video conditioning."), + InputParam( + name="condition_frame_indexes_vision", + type_hint=tuple[int, ...] | list[int], + default=(0, 1), + description="Latent-frame indexes to preserve from the conditioning video.", + ), + InputParam( + name="condition_video_keep", + type_hint=str, + default="first", + description="Which end of a longer conditioning video to use: `first` or `last`.", + ), + InputParam(name="num_frames", type_hint=int, required=True, description="Number of frames to generate."), + InputParam( + name="height", type_hint=int, required=True, description="Height of the generated video in pixels." + ), + InputParam( + name="width", type_hint=int, required=True, description="Width of the generated video in pixels." + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "x0_tokens_vision", + type_hint=torch.Tensor, + description="Vision latents encoded from the conditioning image or video.", + ), + OutputParam( + "vision_condition_frames", + type_hint=list[int], + description="Latent-frame indexes fixed by visual conditioning.", + ), + ] + + @torch.no_grad() + def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + device = components._execution_device + dtype = components.transformer.dtype + + if block_state.video is None: + raise ValueError("`Cosmos3VideoVaeEncoderStep` requires a `video` input.") + if block_state.num_frames == 1: + raise ValueError("`video` conditioning requires `num_frames` > 1.") + if block_state.num_frames < 1: + raise ValueError(f"`num_frames` must be >= 1, got {block_state.num_frames}.") + + sf = int(components.vae.config.scale_factor_spatial) + if block_state.height % sf != 0 or block_state.width % sf != 0: + raise ValueError( + f"`height` and `width` must be multiples of {sf}, got ({block_state.height}, {block_state.width})." + ) + + if not isinstance(block_state.condition_frame_indexes_vision, (list, tuple)) or isinstance( + block_state.condition_frame_indexes_vision, (str, bytes) + ): + raise ValueError( + "`condition_frame_indexes_vision` must be a list/tuple of non-negative ints, e.g. [0, 1]; got " + f"{block_state.condition_frame_indexes_vision!r}." + ) + if not all(isinstance(index, int) and index >= 0 for index in block_state.condition_frame_indexes_vision): + raise ValueError( + "`condition_frame_indexes_vision` must be a list/tuple of non-negative ints, e.g. [0, 1]; got " + f"{block_state.condition_frame_indexes_vision!r}." + ) + if block_state.condition_video_keep not in {"first", "last"}: + raise ValueError("`condition_video_keep` must be either 'first' or 'last'.") + + indexes = tuple(block_state.condition_frame_indexes_vision) + if not indexes: + raise ValueError("`condition_frame_indexes_vision` must contain at least one index.") + latent_t = (block_state.num_frames - 1) // int(components.vae.config.scale_factor_temporal) + 1 + if max(indexes) >= latent_t: + raise ValueError( + f"`condition_frame_indexes_vision` {indexes} contains an index outside the latent timeline " + f"(latent_frames={latent_t} for num_frames={block_state.num_frames})." + ) + + condition_indexes_vision = indexes + conditioning_frames_3d = components.video_processor.preprocess_video( + block_state.video, height=block_state.height, width=block_state.width + ).to(device=device, dtype=dtype) + temporal_compression = int(components.vae.config.scale_factor_temporal) + max_cond_frames = max(condition_indexes_vision) * temporal_compression + 1 + if block_state.condition_video_keep == "first": + conditioning_frames_3d = conditioning_frames_3d[:, :, :max_cond_frames] + else: + conditioning_frames_3d = conditioning_frames_3d[:, :, -max_cond_frames:] + + vision_tensor = torch.zeros( + 1, + 3, + block_state.num_frames, + block_state.height, + block_state.width, + dtype=dtype, + device=device, + ) + t_fill = min(conditioning_frames_3d.shape[2], block_state.num_frames) + vision_tensor[:, :, :t_fill] = conditioning_frames_3d[:, :, :t_fill] + if t_fill < block_state.num_frames: + vision_tensor[:, :, t_fill:] = vision_tensor[:, :, t_fill - 1 : t_fill].expand( + -1, -1, block_state.num_frames - t_fill, -1, -1 + ) + vision_condition_frames = list(condition_indexes_vision) + + block_state.x0_tokens_vision = components._encode_video(vision_tensor).contiguous().float() + block_state.vision_condition_frames = vision_condition_frames + + self.set_block_state(state, block_state) + return components, state + + +class Cosmos3TransferChunkVaeEncoderStep(ModularPipelineBlocks): + model_name = "cosmos3-omni" + + @property + def description(self) -> str: + return ( + "Per-chunk transfer VAE encode: slices + pads this chunk's control maps, seeds the target's conditioning " + "frames (first chunk from the input video, later chunks from the previous chunk's tail), and encodes both " + "the controls and the seeded target into clean Cosmos3 vision latents. Runs inside the autoregressive " + "chunk loop." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("vae", AutoencoderKLWan), + ComponentSpec( + "video_processor", + VideoProcessor, + config=FrozenDict({"vae_scale_factor": 16, "resample": "bilinear"}), + default_creation_method="from_config", + ), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam(name="chunk_id", type_hint=int, default=0, description="Index of the current chunk."), + InputParam( + name="previous_output", + default=None, + description="Decoded pixels of the previous chunk, used to seed later chunks.", + ), + InputParam( + name="control_frames", + type_hint=dict, + required=True, + description="Preprocessed, time-padded control maps in canonical hint order.", + ), + InputParam(name="chunk_frames", type_hint=int, required=True, description="Pixel frames per chunk."), + InputParam( + name="total_frames", type_hint=int, required=True, description="Total number of output frames." + ), + InputParam(name="stride", type_hint=int, required=True, description="Frame stride between chunks."), + InputParam( + name="height", type_hint=int, required=True, description="Height of the generated video in pixels." + ), + InputParam( + name="width", type_hint=int, required=True, description="Width of the generated video in pixels." + ), + InputParam( + name="video", + default=None, + description="Optional input video that seeds the first chunk's conditioning.", + ), + InputParam( + name="num_first_chunk_conditional_frames", + type_hint=int, + default=0, + description="Number of frames the first chunk reuses from the input video.", + ), + InputParam( + name="num_conditional_frames", + type_hint=int, + default=1, + description="Number of frames each later chunk reuses from the previous chunk's tail.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "control_latents", + type_hint=list[torch.Tensor], + description="Clean control latents for this chunk, one per hint in canonical order.", + ), + OutputParam( + "x0_tokens_vision", + type_hint=torch.Tensor, + description="Clean target vision latents encoded from the seeded target frames.", + ), + OutputParam( + "current_conditional_frames", + type_hint=int, + description="Number of pixel frames actually used to seed this chunk's target.", + ), + ] + + @torch.no_grad() + def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + device = components._execution_device + dtype = components.transformer.dtype + + chunk_id = block_state.chunk_id + chunk_frames = block_state.chunk_frames + height = block_state.height + width = block_state.width + + # Slice this chunk's window out of the (padded) control maps and reflect-pad it up to a full chunk (repeat the + # last frame once too short to keep reflecting). control_frames is already in canonical hint order. + start_frame = chunk_id * block_state.stride + end_frame = min(start_frame + chunk_frames, block_state.total_frames) + chunk_controls = [] + for frames in block_state.control_frames.values(): + frames = frames[:, :, start_frame:end_frame] + while frames.shape[2] < chunk_frames: + pad_len = min(frames.shape[2] - 1, chunk_frames - frames.shape[2]) + if pad_len <= 0: + pad_frame = frames[:, :, -1:].repeat(1, 1, chunk_frames - frames.shape[2], 1, 1) + frames = torch.cat([frames, pad_frame], dim=2) + break + frames = torch.cat([frames, frames.flip(dims=[2])[:, :, :pad_len]], dim=2) + chunk_controls.append(frames) + + # Seed the target with conditioning frames (first chunk from the input video, later chunks from the + # previous chunk's tail), repeat-padding the remaining frames so the whole clip is well-defined. + target = torch.zeros(1, 3, chunk_frames, height, width, device=device, dtype=dtype) + current_conditional_frames = 0 + if chunk_id == 0 and block_state.num_first_chunk_conditional_frames > 0 and block_state.video is not None: + input_frames = components.video_processor.preprocess_video( + block_state.video, height=height, width=width + ).to(device=device, dtype=dtype) + current_conditional_frames = min( + block_state.num_first_chunk_conditional_frames, input_frames.shape[2], chunk_frames + ) + if current_conditional_frames > 0: + target[:, :, :current_conditional_frames] = input_frames[:, :, :current_conditional_frames] + elif chunk_id > 0 and block_state.previous_output is not None: + current_conditional_frames = min( + block_state.num_conditional_frames, block_state.previous_output.shape[2], chunk_frames + ) + if current_conditional_frames > 0: + target[:, :, :current_conditional_frames] = block_state.previous_output[ + :, :, -current_conditional_frames: + ].to(device=device, dtype=dtype) + if 0 < current_conditional_frames < chunk_frames: + fill = target[:, :, current_conditional_frames - 1 : current_conditional_frames] + target[:, :, current_conditional_frames:] = fill.expand( + -1, -1, chunk_frames - current_conditional_frames, -1, -1 + ) + + block_state.control_latents = [components._encode_video(ctrl).contiguous().float() for ctrl in chunk_controls] + block_state.x0_tokens_vision = components._encode_video(target).contiguous().float() + block_state.current_conditional_frames = current_conditional_frames + + self.set_block_state(state, block_state) + return components, state + + +class Cosmos3ActionVisionVaeEncoderStep(ModularPipelineBlocks): + model_name = "cosmos3-omni" + + @property + def description(self) -> str: + return ( + "Prepares action-conditioned vision latents and action frame metadata. " + "Only the action visual reference (image/video) is VAE-encoded; action vectors are handled separately." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("vae", AutoencoderKLWan), + ComponentSpec( + "video_processor", + VideoProcessor, + config=FrozenDict({"vae_scale_factor": 16, "resample": "bilinear"}), + default_creation_method="from_config", + ), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + name="action", + type_hint=CosmosActionCondition, + required=True, + description="Action-conditioning metadata and its reference visual input.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "x0_tokens_vision", + type_hint=torch.Tensor, + description="Vision latents encoded from the conditioning image or video.", + ), + OutputParam( + "vision_condition_frames", + type_hint=list[int], + description="Latent-frame indexes fixed by visual conditioning.", + ), + OutputParam( + "action_condition_frame_indexes", + type_hint=list[int], + description="Action-frame indexes fixed by action conditioning.", + ), + ] + + @torch.no_grad() + def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + device = components._execution_device + dtype = components.vae.dtype + + action = block_state.action + target_frames = action.chunk_size + 1 + conditioning_clip = [action.image] if action.image is not None else action.video + vision_tensor, action_image_size, _, _ = components._prepare_action_video_conditioning( + conditioning_clip, + action.resolution_tier, + target_frames, + device=device, + dtype=dtype, + ) + + if action.mode == "forward_dynamics": + vision_condition_frames = [0] + action_condition_frame_indexes = list(range(action.chunk_size)) + elif action.mode == "policy": + vision_condition_frames = [0] + action_condition_frame_indexes = [] + elif action.mode == "inverse_dynamics": + latent_frames = (target_frames - 1) // int(components.vae.config.scale_factor_temporal) + 1 + vision_condition_frames = list(range(latent_frames)) + action_condition_frame_indexes = [] + else: + raise ValueError( + f"Unsupported action_mode={action.mode!r}; expected one of ['forward_dynamics', 'inverse_dynamics', 'policy']." + ) + + x0_tokens_vision = components._encode_video(vision_tensor).contiguous().float() + if action_image_size is not None: + x0_tokens_vision = components._remove_action_video_padding_from_latent(x0_tokens_vision, action_image_size) + + block_state.x0_tokens_vision = x0_tokens_vision + block_state.vision_condition_frames = vision_condition_frames + block_state.action_condition_frame_indexes = action_condition_frame_indexes + + self.set_block_state(state, block_state) + return components, state diff --git a/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3.py b/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3.py new file mode 100644 index 000000000000..726eb656f71b --- /dev/null +++ b/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3.py @@ -0,0 +1,1039 @@ +import torch + +from ..modular_pipeline import ( + AutoPipelineBlocks, + ConditionalPipelineBlocks, + PipelineState, + SequentialPipelineBlocks, +) +from ..modular_pipeline_utils import InputParam, OutputParam +from .after_decode import Cosmos3ActionOutputStep +from .before_denoise import ( + Cosmos3ActionDenoiseInputStep, + Cosmos3ActionPackSequenceStep, + Cosmos3ActionPrepareLatentsStep, + Cosmos3PrepareTextSegmentsStep, + Cosmos3SetTimestepsStep, + Cosmos3SoundDenoiseInputStep, + Cosmos3SoundPackSequenceStep, + Cosmos3SoundPrepareLatentsStep, + Cosmos3TransferPackSequenceStep, + Cosmos3TransferPrepareLatentsStep, + Cosmos3TransferSetTimestepsStep, + Cosmos3TransferSetupStep, + Cosmos3VisionDenoiseInputStep, + Cosmos3VisionPackSequenceStep, + Cosmos3VisionPrepareLatentsStep, +) +from .decoders import ( + Cosmos3SoundDecodeStep, + Cosmos3TransferDecodeChunkStep, + Cosmos3TransferStitchStep, + Cosmos3VideoDecodeStep, +) +from .denoise import ( + Cosmos3TransferDenoiseStep, + Cosmos3VisionActionDenoiseStep, + Cosmos3VisionDenoiseStep, + Cosmos3VisionSoundActionDenoiseStep, + Cosmos3VisionSoundDenoiseStep, +) +from .encoders import ( + Cosmos3ActionTextStep, + Cosmos3ActionVisionVaeEncoderStep, + Cosmos3ImageVaeEncoderStep, + Cosmos3TextEncoderStep, + Cosmos3TransferChunkVaeEncoderStep, + Cosmos3TransferTextStep, + Cosmos3VideoVaeEncoderStep, +) +from .modular_pipeline import Cosmos3OmniModularPipeline + + +class Cosmos3TransferTextBlocks(SequentialPipelineBlocks): + model_name = "cosmos3-omni" + block_classes = [Cosmos3TransferSetupStep, Cosmos3TransferTextStep] + block_names = ["setup", "transfer_text"] + + @property + def description(self): + return ( + "Transfer text branch: resolves the control-video chunk geometry, then tokenizes the (pre-upsampled) " + "prompt in transfer mode using the per-chunk frame count." + ) + + +# auto_docstring +class Cosmos3AutoTextEncoderStep(AutoPipelineBlocks): + """ + Auto text encoder block for Cosmos3. + - Cosmos3TransferTextBlocks runs when control_videos are provided. + - Cosmos3ActionTextStep runs when action is provided. + - Cosmos3TextEncoderStep runs otherwise. + + Components: + video_processor (`VideoProcessor`) + text_tokenizer (`AutoTokenizer`) + + Inputs: + control_videos (`dict`, *optional*): + Mapping of hint name (edge/blur/depth/seg/wsm) to the control video for that modality. + height (`int`, *optional*): + Height of the generated video in pixels. + width (`int`, *optional*): + Width of the generated video in pixels. + num_frames (`int`, *optional*): + Optional cap on the number of output frames (defaults to the control video length). + num_video_frames_per_chunk (`int`, *optional*): + Number of pixel frames generated per autoregressive chunk. + num_conditional_frames (`int`, *optional*, defaults to 1): + Number of frames each chunk reuses from the previous chunk's tail. + prompt (`str`): + The text prompt that guides Cosmos3 generation. + negative_prompt (`str`, *optional*): + The negative text prompt used for classifier-free guidance. + fps (`float`, *optional*, defaults to 24.0): + Frame rate of the generated video. + use_system_prompt (`bool`, *optional*, defaults to True): + Whether to prepend the Cosmos3 system prompt. + add_resolution_template (`bool`, *optional*, defaults to True): + Whether to add resolution metadata to the prompt. + add_duration_template (`bool`, *optional*, defaults to True): + Whether to add duration metadata to the prompt. + action (`CosmosActionCondition`, *optional*): + Action-conditioning metadata and its reference visual input. + + Outputs: + height (`int`): + Resolved output height in pixels. + width (`int`): + Resolved output width in pixels. + control_frames (`dict`): + Preprocessed, time-padded control maps in canonical hint order. + total_frames (`int`): + Total number of output frames to generate. + chunk_frames (`int`): + Number of pixel frames per autoregressive chunk. + num_chunks (`int`): + Number of autoregressive chunks. + stride (`int`): + Frame stride between consecutive chunks. + cond_input_ids (`Tensor`): + Token IDs for the conditional prompt. + uncond_input_ids (`Tensor`): + Token IDs for the unconditional prompt. + action_mode (`str`): + Requested action-generation mode. + num_frames (`int`): + Number of frames to generate. + """ + + model_name = "cosmos3-omni" + block_classes = [Cosmos3TransferTextBlocks, Cosmos3ActionTextStep, Cosmos3TextEncoderStep] + block_names = ["transfer_text", "action_text", "text"] + block_trigger_inputs = ["control_videos", "action", None] + + @property + def description(self): + return ( + "Auto text encoder block for Cosmos3.\n" + + " - Cosmos3TransferTextBlocks runs when control_videos are provided.\n" + + " - Cosmos3ActionTextStep runs when action is provided.\n" + + " - Cosmos3TextEncoderStep runs otherwise." + ) + + +# auto_docstring +class Cosmos3AutoVaeEncoderStep(ConditionalPipelineBlocks): + """ + Auto VAE conditioning block for Cosmos3. + - Cosmos3ActionVisionVaeEncoderStep runs when action is provided. + - Cosmos3VideoVaeEncoderStep runs for the non-action video path. + - Cosmos3ImageVaeEncoderStep runs for the non-action image path. + - when no action, image, or video conditioning is provided, this block is skipped. + + Components: + vae (`AutoencoderKLWan`) + video_processor (`VideoProcessor`) + + Inputs: + action (`CosmosActionCondition`, *optional*): + Action-conditioning metadata and its reference visual input. + video (`None`, *optional*): + Reference video for video-to-video conditioning. + condition_frame_indexes_vision (`tuple | list`, *optional*, defaults to (0, 1)): + Latent-frame indexes to preserve from the conditioning video. + condition_video_keep (`str`, *optional*, defaults to first): + Which end of a longer conditioning video to use: `first` or `last`. + num_frames (`int`, *optional*): + Number of frames to generate. + height (`int`, *optional*): + Height of the generated video in pixels. + width (`int`, *optional*): + Width of the generated video in pixels. + image (`None`, *optional*): + Reference image for image-to-video conditioning. + + Outputs: + x0_tokens_vision (`Tensor`): + Vision latents encoded from the conditioning image or video. + vision_condition_frames (`list`): + Latent-frame indexes fixed by visual conditioning. + action_condition_frame_indexes (`list`): + Action-frame indexes fixed by action conditioning. + """ + + model_name = "cosmos3-omni" + block_classes = [Cosmos3ActionVisionVaeEncoderStep, Cosmos3VideoVaeEncoderStep, Cosmos3ImageVaeEncoderStep] + block_names = ["action_conditioning", "video_conditioning", "image_conditioning"] + block_trigger_inputs = ["action", "video", "image", "control_videos"] + default_block_name = None + + def select_block(self, **kwargs) -> str | None: + action = kwargs.get("action") + image = kwargs.get("image") + video = kwargs.get("video") + # Transfer preprocesses/encodes its control maps inside the denoise chunk loop, so the standard VAE + # conditioning stage is skipped when control_videos drive the workflow. + if kwargs.get("control_videos") is not None: + return None + if action is not None: + if image is not None or video is not None: + raise ValueError( + "Pass action conditioning via `action.image` / `action.video`, not top-level image/video." + ) + return "action_conditioning" + if image is not None and video is not None: + raise ValueError("Pass either image or video, not both.") + if video is not None: + return "video_conditioning" + if image is not None: + return "image_conditioning" + return None + + @property + def description(self): + return ( + "Auto VAE conditioning block for Cosmos3.\n" + + " - Cosmos3ActionVisionVaeEncoderStep runs when action is provided.\n" + + " - Cosmos3VideoVaeEncoderStep runs for the non-action video path.\n" + + " - Cosmos3ImageVaeEncoderStep runs for the non-action image path.\n" + + " - when no action, image, or video conditioning is provided, this block is skipped." + ) + + +# auto_docstring +class Cosmos3AutoSoundDecodeStep(AutoPipelineBlocks): + """ + Auto sound decoder block for Cosmos3. + - Cosmos3SoundDecodeStep runs when sound_latents are present. + - if sound_latents are not provided, this block is skipped. + + Components: + sound_tokenizer (`Cosmos3AVAEAudioTokenizer`) + + Inputs: + sound_latents (`Tensor`, *optional*): + Denoised sound latents to decode. + + Outputs: + sound (`Tensor`): + Generated waveform. + sampling_rate (`int`): + Sample rate of the generated waveform in Hz. + """ + + model_name = "cosmos3-omni" + block_classes = [Cosmos3SoundDecodeStep] + block_names = ["decode"] + block_trigger_inputs = ["sound_latents"] + + @property + def description(self): + return ( + "Auto sound decoder block for Cosmos3.\n" + + " - Cosmos3SoundDecodeStep runs when sound_latents are present.\n" + + " - if sound_latents are not provided, this block is skipped." + ) + + +# auto_docstring +class Cosmos3DecodeStep(SequentialPipelineBlocks): + """ + Decodes denoised latents into modality outputs. + + Components: + vae (`AutoencoderKLWan`) + video_processor (`VideoProcessor`) + sound_tokenizer (`Cosmos3AVAEAudioTokenizer`) + + Inputs: + latents (`Tensor`): + Denoised vision latents to decode. + output_type (`str`, *optional*, defaults to pil): + Output format: 'pil', 'np', 'pt'. + sound_latents (`Tensor`, *optional*): + Denoised sound latents to decode. + + Outputs: + videos (`list`): + The generated videos. + sound (`Tensor`): + Generated waveform. + sampling_rate (`int`): + Sample rate of the generated waveform in Hz. + """ + + model_name = "cosmos3-omni" + block_classes = [Cosmos3VideoDecodeStep, Cosmos3AutoSoundDecodeStep] + block_names = ["video", "sound"] + + @property + def description(self) -> str: + return "Decodes denoised latents into modality outputs." + + +class Cosmos3AutoDecodeStep(ConditionalPipelineBlocks): + model_name = "cosmos3-omni" + block_classes = [Cosmos3TransferStitchStep, Cosmos3DecodeStep] + block_names = ["transfer", "standard"] + block_trigger_inputs = ["control_videos"] + default_block_name = "standard" + + def select_block(self, **kwargs) -> str | None: + if kwargs.get("control_videos") is not None: + return "transfer" + return "standard" + + @property + def description(self) -> str: + return ( + "Selects the Cosmos3 decode workflow.\n" + + " - Cosmos3TransferStitchStep stitches the decoded transfer chunks when control_videos are provided.\n" + + " - Cosmos3DecodeStep decodes the denoised latents otherwise." + ) + + +# auto_docstring +class Cosmos3VisionCoreDenoiseStep(SequentialPipelineBlocks): + """ + Runs the text-and-vision Cosmos3 denoising workflow. + + Components: + transformer (`Cosmos3OmniTransformer`) + scheduler (`UniPCMultistepScheduler`) + + Inputs: + cond_input_ids (`None`): + Token IDs for the conditional prompt. + uncond_input_ids (`None`): + Token IDs for the unconditional prompt. + x0_tokens_vision (`Tensor`, *optional*): + Vision latents encoded from the conditioning image or video. + vision_condition_frames (`list`, *optional*): + Latent-frame indexes fixed by visual conditioning. + num_frames (`int`): + Number of frames to generate. + height (`int`): + Height of the generated video in pixels. + width (`int`): + Width of the generated video in pixels. + fps (`float`, *optional*, defaults to 24.0): + Frame rate of the generated video. + latents (`Tensor`, *optional*): + Pre-generated noisy vision latents. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + num_inference_steps (`int`): + The number of denoising steps. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + guidance_scale (`float`, *optional*, defaults to 6.0): + Scale for classifier-free guidance. + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + model_name = "cosmos3-omni" + block_classes = [ + Cosmos3PrepareTextSegmentsStep, + Cosmos3VisionPrepareLatentsStep, + Cosmos3VisionPackSequenceStep, + Cosmos3VisionDenoiseInputStep, + Cosmos3SetTimestepsStep, + Cosmos3VisionDenoiseStep, + ] + block_names = [ + "prepare_text_segments", + "prepare_vision_latents", + "pack_vision_sequence", + "prepare_vision_denoiser_inputs", + "set_timesteps", + "denoise", + ] + + @property + def description(self): + return "Runs the text-and-vision Cosmos3 denoising workflow." + + @property + def outputs(self): + return [OutputParam.template("latents")] + + +# auto_docstring +class Cosmos3VisionSoundCoreDenoiseStep(SequentialPipelineBlocks): + """ + Runs the text, vision, and sound Cosmos3 denoising workflow. + + Components: + transformer (`Cosmos3OmniTransformer`) + scheduler (`UniPCMultistepScheduler`) + + Inputs: + cond_input_ids (`None`): + Token IDs for the conditional prompt. + uncond_input_ids (`None`): + Token IDs for the unconditional prompt. + x0_tokens_vision (`Tensor`, *optional*): + Vision latents encoded from the conditioning image or video. + vision_condition_frames (`list`, *optional*): + Latent-frame indexes fixed by visual conditioning. + num_frames (`int`): + Number of frames to generate. + height (`int`): + Height of the generated video in pixels. + width (`int`): + Width of the generated video in pixels. + fps (`float`, *optional*, defaults to 24.0): + Frame rate of the generated video. + latents (`Tensor`, *optional*): + Pre-generated noisy vision latents. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + num_inference_steps (`int`): + The number of denoising steps. + sound_latents (`Tensor`, *optional*): + Pre-generated noisy sound latents. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + guidance_scale (`float`, *optional*, defaults to 6.0): + Scale for classifier-free guidance. + + Outputs: + latents (`Tensor`): + Denoised latents. + sound_latents (`Tensor`): + Denoised sound latents. + """ + + model_name = "cosmos3-omni" + block_classes = [ + Cosmos3PrepareTextSegmentsStep, + Cosmos3VisionPrepareLatentsStep, + Cosmos3VisionPackSequenceStep, + Cosmos3VisionDenoiseInputStep, + Cosmos3SetTimestepsStep, + Cosmos3SoundPrepareLatentsStep, + Cosmos3SoundPackSequenceStep, + Cosmos3SoundDenoiseInputStep, + Cosmos3VisionSoundDenoiseStep, + ] + block_names = [ + "prepare_text_segments", + "prepare_vision_latents", + "pack_vision_sequence", + "prepare_vision_denoiser_inputs", + "set_timesteps", + "prepare_sound_latents", + "pack_sound_sequence", + "prepare_sound_denoiser_inputs", + "denoise", + ] + + @property + def description(self): + return "Runs the text, vision, and sound Cosmos3 denoising workflow." + + @property + def outputs(self): + return [ + OutputParam.template("latents"), + OutputParam("sound_latents", type_hint=torch.Tensor, description="Denoised sound latents."), + ] + + +# auto_docstring +class Cosmos3VisionActionCoreDenoiseStep(SequentialPipelineBlocks): + """ + Runs the text, vision, and action Cosmos3 denoising workflow. + + Components: + transformer (`Cosmos3OmniTransformer`) + scheduler (`UniPCMultistepScheduler`) + + Inputs: + cond_input_ids (`None`): + Token IDs for the conditional prompt. + uncond_input_ids (`None`): + Token IDs for the unconditional prompt. + x0_tokens_vision (`Tensor`, *optional*): + Vision latents encoded from the conditioning image or video. + vision_condition_frames (`list`, *optional*): + Latent-frame indexes fixed by visual conditioning. + num_frames (`int`): + Number of frames to generate. + height (`int`): + Height of the generated video in pixels. + width (`int`): + Width of the generated video in pixels. + fps (`float`, *optional*, defaults to 24.0): + Frame rate of the generated video. + latents (`Tensor`, *optional*): + Pre-generated noisy vision latents. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + num_inference_steps (`int`): + The number of denoising steps. + action (`CosmosActionCondition`): + Action-conditioning metadata. + action_condition_frame_indexes (`list`, *optional*): + Action-frame indexes fixed by action conditioning. + action_latents (`Tensor`, *optional*): + Pre-generated noisy action latents. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + guidance_scale (`float`, *optional*, defaults to 6.0): + Scale for classifier-free guidance. + + Outputs: + latents (`Tensor`): + Denoised latents. + action_latents (`Tensor`): + Denoised action latents. + """ + + model_name = "cosmos3-omni" + block_classes = [ + Cosmos3PrepareTextSegmentsStep, + Cosmos3VisionPrepareLatentsStep, + Cosmos3VisionPackSequenceStep, + Cosmos3VisionDenoiseInputStep, + Cosmos3SetTimestepsStep, + Cosmos3ActionPrepareLatentsStep, + Cosmos3ActionPackSequenceStep, + Cosmos3ActionDenoiseInputStep, + Cosmos3VisionActionDenoiseStep, + ] + block_names = [ + "prepare_text_segments", + "prepare_vision_latents", + "pack_vision_sequence", + "prepare_vision_denoiser_inputs", + "set_timesteps", + "prepare_action_latents", + "pack_action_sequence", + "prepare_action_denoiser_inputs", + "denoise", + ] + + @property + def description(self): + return "Runs the text, vision, and action Cosmos3 denoising workflow." + + @property + def outputs(self): + return [ + OutputParam.template("latents"), + OutputParam("action_latents", type_hint=torch.Tensor, description="Denoised action latents."), + ] + + +# auto_docstring +class Cosmos3VisionSoundActionCoreDenoiseStep(SequentialPipelineBlocks): + """ + Runs the text, vision, sound, and action Cosmos3 denoising workflow. + + Components: + transformer (`Cosmos3OmniTransformer`) + scheduler (`UniPCMultistepScheduler`) + + Inputs: + cond_input_ids (`None`): + Token IDs for the conditional prompt. + uncond_input_ids (`None`): + Token IDs for the unconditional prompt. + x0_tokens_vision (`Tensor`, *optional*): + Vision latents encoded from the conditioning image or video. + vision_condition_frames (`list`, *optional*): + Latent-frame indexes fixed by visual conditioning. + num_frames (`int`): + Number of frames to generate. + height (`int`): + Height of the generated video in pixels. + width (`int`): + Width of the generated video in pixels. + fps (`float`, *optional*, defaults to 24.0): + Frame rate of the generated video. + latents (`Tensor`, *optional*): + Pre-generated noisy vision latents. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + num_inference_steps (`int`): + The number of denoising steps. + sound_latents (`Tensor`, *optional*): + Pre-generated noisy sound latents. + action (`CosmosActionCondition`): + Action-conditioning metadata. + action_condition_frame_indexes (`list`, *optional*): + Action-frame indexes fixed by action conditioning. + action_latents (`Tensor`, *optional*): + Pre-generated noisy action latents. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + guidance_scale (`float`, *optional*, defaults to 6.0): + Scale for classifier-free guidance. + + Outputs: + latents (`Tensor`): + Denoised latents. + sound_latents (`Tensor`): + Denoised sound latents. + action_latents (`Tensor`): + Denoised action latents. + """ + + model_name = "cosmos3-omni" + block_classes = [ + Cosmos3PrepareTextSegmentsStep, + Cosmos3VisionPrepareLatentsStep, + Cosmos3VisionPackSequenceStep, + Cosmos3VisionDenoiseInputStep, + Cosmos3SetTimestepsStep, + Cosmos3SoundPrepareLatentsStep, + Cosmos3SoundPackSequenceStep, + Cosmos3SoundDenoiseInputStep, + Cosmos3ActionPrepareLatentsStep, + Cosmos3ActionPackSequenceStep, + Cosmos3ActionDenoiseInputStep, + Cosmos3VisionSoundActionDenoiseStep, + ] + block_names = [ + "prepare_text_segments", + "prepare_vision_latents", + "pack_vision_sequence", + "prepare_vision_denoiser_inputs", + "set_timesteps", + "prepare_sound_latents", + "pack_sound_sequence", + "prepare_sound_denoiser_inputs", + "prepare_action_latents", + "pack_action_sequence", + "prepare_action_denoiser_inputs", + "denoise", + ] + + @property + def description(self): + return "Runs the text, vision, sound, and action Cosmos3 denoising workflow." + + @property + def outputs(self): + return [ + OutputParam.template("latents"), + OutputParam("sound_latents", type_hint=torch.Tensor, description="Denoised sound latents."), + OutputParam("action_latents", type_hint=torch.Tensor, description="Denoised action latents."), + ] + + +class Cosmos3TransferChunkDenoiseStep(SequentialPipelineBlocks): + model_name = "cosmos3-omni" + block_classes = [ + Cosmos3TransferChunkVaeEncoderStep, + Cosmos3TransferPrepareLatentsStep, + Cosmos3TransferPackSequenceStep, + Cosmos3TransferSetTimestepsStep, + Cosmos3TransferDenoiseStep, + Cosmos3TransferDecodeChunkStep, + ] + block_names = [ + "encode_transfer_chunk", + "prepare_transfer_latents", + "pack_transfer_sequence", + "set_timesteps", + "denoise", + "decode_chunk", + ] + + @property + def description(self) -> str: + return ( + "Autoregressive transfer chunk loop. Overrides __call__ to iterate chunks (the inner timestep loop is a " + "non-leaf LoopSequentialPipelineBlocks, so this outer loop cannot itself be a LoopSequentialPipelineBlocks). " + "Per-chunk cross-carry (previous_output, output_chunks) lives on PipelineState." + ) + + @property + def inputs(self) -> list[InputParam]: + return super().inputs + [ + InputParam(name="num_chunks", type_hint=int, required=True, description="Number of autoregressive chunks.") + ] + + @torch.no_grad() + def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState: + num_chunks = state.get("num_chunks") + state.set("output_chunks", []) + state.set("previous_output", None) + for chunk_id in range(num_chunks): + state.set("chunk_id", chunk_id) + for _, block in self.sub_blocks.items(): + components, state = block(components, state) + return components, state + + +class Cosmos3TransferCoreDenoiseStep(SequentialPipelineBlocks): + model_name = "cosmos3-omni" + block_classes = [ + Cosmos3PrepareTextSegmentsStep, + Cosmos3TransferChunkDenoiseStep, + ] + block_names = ["prepare_text_segments", "chunk_denoise"] + + @property + def description(self) -> str: + return "Transfer denoise stage: prepare shared text segments once, then run the autoregressive chunk loop." + + +# auto_docstring +class Cosmos3AutoCoreDenoiseStep(ConditionalPipelineBlocks): + """ + Selects the Cosmos3 core denoising workflow. + - transfer runs the autoregressive control-video (ControlNet-style) chunk loop when control_videos are provided. + - vision_sound_action runs when action and enable_sound are provided. + - vision_action runs when action is provided. + - vision_sound runs when enable_sound is true. + - vision runs otherwise. + + Components: + transformer (`Cosmos3OmniTransformer`) + vae (`AutoencoderKLWan`) + video_processor (`VideoProcessor`) + scheduler (`UniPCMultistepScheduler`) + + Inputs: + cond_input_ids (`None`): + Token IDs for the conditional prompt. + uncond_input_ids (`None`): + Token IDs for the unconditional prompt. + chunk_id (`int`, *optional*, defaults to 0): + Index of the current chunk. + previous_output (`None`, *optional*): + Decoded pixels of the previous chunk, used to seed later chunks. + control_frames (`dict`, *optional*): + Preprocessed, time-padded control maps in canonical hint order. + chunk_frames (`int`, *optional*): + Pixel frames per chunk. + total_frames (`int`, *optional*): + Total number of output frames. + stride (`int`, *optional*): + Frame stride between chunks. + height (`int`): + Height of the generated video in pixels. + width (`int`): + Width of the generated video in pixels. + video (`None`, *optional*): + Optional input video that seeds the first chunk's conditioning. + num_first_chunk_conditional_frames (`int`, *optional*, defaults to 0): + Number of frames the first chunk reuses from the input video. + num_conditional_frames (`int`, *optional*, defaults to 1): + Number of frames each later chunk reuses from the previous chunk's tail. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + fps (`float`, *optional*, defaults to 24.0): + Frame rate of the generated video. + share_vision_temporal_positions (`bool`, *optional*, defaults to True): + Whether control and target items share vision temporal positions. + num_inference_steps (`int`): + The number of denoising steps. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + guidance_scale (`float`, *optional*, defaults to 6.0): + Scale for text classifier-free guidance. + control_guidance (`float`, *optional*, defaults to 1.0): + Scale for the control (structural) guidance axis. + guidance_interval (`tuple`, *optional*): + Timestep interval [lo, hi] over which text guidance is active (None = always). + control_guidance_interval (`tuple`, *optional*): + Timestep interval [lo, hi] over which control guidance is active (None = always). + output_chunks (`list`, *optional*): + Decoded pixel chunks accumulated so far. + num_chunks (`int`, *optional*): + Number of autoregressive chunks. + x0_tokens_vision (`Tensor`, *optional*): + Vision latents encoded from the conditioning image or video. + vision_condition_frames (`list`, *optional*): + Latent-frame indexes fixed by visual conditioning. + num_frames (`int`, *optional*): + Number of frames to generate. + latents (`Tensor`): + Pre-generated noisy vision latents. + sound_latents (`Tensor`, *optional*): + Pre-generated noisy sound latents. + action (`CosmosActionCondition`, *optional*): + Action-conditioning metadata. + action_condition_frame_indexes (`list`, *optional*): + Action-frame indexes fixed by action conditioning. + action_latents (`Tensor`, *optional*): + Pre-generated noisy action latents. + enable_sound (`bool`, *optional*, defaults to False): + Whether to generate a synchronized sound track. + + Outputs: + cond_text_segment (`dict`): + Conditional text segment for the denoiser. + uncond_text_segment (`dict`): + Unconditional text segment for the denoiser. + control_latents (`list`): + Clean control latents for this chunk, one per hint in canonical order. + x0_tokens_vision (`Tensor`): + Clean target vision latents encoded from the seeded target frames. + current_conditional_frames (`int`): + Number of pixel frames actually used to seed this chunk's target. + latents (`Tensor`): + Noisy target latents for this chunk. + velocity_mask (`Tensor`): + Mask that zeroes the velocity on conditioned (clean) latent frames. + condition_latents (`Tensor`): + Clean target latents on the conditioned frames (the autoregressive seed). + target_condition_indexes (`list`): + Latent-frame indexes fixed by the chunk's conditioning. + cond_full_static (`dict`): + Conditional [control..., target] transfer sequence carrying every control item. + cond_no_control_static (`dict`): + Conditional [target] transfer sequence with the control items dropped. + uncond_full_static (`dict`): + Unconditional [control..., target] transfer sequence for text CFG. + num_noisy_vision_tokens (`int`): + Number of noisy target vision tokens denoised each step. + timesteps (`Tensor`): + Scheduler timesteps for this chunk. + num_warmup_steps (`int`): + Number of scheduler warmup steps for this chunk. + vision_tokens_full (`list`): + Token list for the [control..., target] forward passes. + vision_tokens_target (`list`): + Token list for the target-only (no-control) forward pass. + vision_timesteps (`Tensor`): + Timesteps for the noisy target tokens. + velocity (`Tensor`): + Predicted (masked) transfer velocity. + previous_output (`Tensor`): + Decoded pixels of this chunk, used to seed the next chunk. + output_chunks (`list`): + Decoded pixel chunks accumulated so far (with this chunk appended). + sound_latents (`Tensor`): + Denoised sound latents. + action_latents (`Tensor`): + Denoised action latents. + """ + + model_name = "cosmos3-omni" + block_classes = [ + Cosmos3TransferCoreDenoiseStep, + Cosmos3VisionSoundActionCoreDenoiseStep, + Cosmos3VisionActionCoreDenoiseStep, + Cosmos3VisionSoundCoreDenoiseStep, + Cosmos3VisionCoreDenoiseStep, + ] + block_names = ["transfer", "vision_sound_action", "vision_action", "vision_sound", "vision"] + block_trigger_inputs = ["action", "enable_sound", "control_videos"] + default_block_name = "vision" + + @property + def inputs(self): + inputs = super().inputs + inputs.append( + InputParam( + name="enable_sound", + type_hint=bool, + default=False, + description="Whether to generate a synchronized sound track.", + ) + ) + return inputs + + def select_block(self, **kwargs) -> str | None: + action = kwargs.get("action") + enable_sound = kwargs.get("enable_sound") + if kwargs.get("control_videos") is not None: + return "transfer" + if action is not None and enable_sound: + return "vision_sound_action" + if action is not None: + return "vision_action" + if enable_sound: + return "vision_sound" + return "vision" + + @property + def description(self): + return ( + "Selects the Cosmos3 core denoising workflow.\n" + + " - transfer runs the autoregressive control-video (ControlNet-style) chunk loop when control_videos are provided.\n" + + " - vision_sound_action runs when action and enable_sound are provided.\n" + + " - vision_action runs when action is provided.\n" + + " - vision_sound runs when enable_sound is true.\n" + + " - vision runs otherwise." + ) + + +# auto_docstring +class Cosmos3OmniBlocks(SequentialPipelineBlocks): + """ + Modular pipeline blocks for Cosmos3 generation modes. + + Supported workflows: + - `text2image`: requires `prompt`, `num_frames` + - `text2video`: requires `prompt` + - `image2video`: requires `prompt`, `image` + - `video2video`: requires `prompt`, `video` + - `text2video_with_sound`: requires `prompt`, `enable_sound` + - `image2video_with_sound`: requires `prompt`, `image`, `enable_sound` + - `video2video_with_sound`: requires `prompt`, `video`, `enable_sound` + - `action_policy`: requires `prompt`, `action` + - `action_forward_dynamics`: requires `prompt`, `action` + - `action_inverse_dynamics`: requires `prompt`, `action` + - `transfer`: requires `prompt`, `control_videos` + + Components: + video_processor (`VideoProcessor`) + text_tokenizer (`AutoTokenizer`) + vae (`AutoencoderKLWan`) + transformer (`Cosmos3OmniTransformer`) + scheduler (`UniPCMultistepScheduler`) + sound_tokenizer (`Cosmos3AVAEAudioTokenizer`) + + Inputs: + control_videos (`dict`, *optional*): + Mapping of hint name (edge/blur/depth/seg/wsm) to the control video for that modality. + height (`int`, *optional*): + Height of the generated video in pixels. + width (`int`, *optional*): + Width of the generated video in pixels. + num_frames (`int`, *optional*): + Optional cap on the number of output frames (defaults to the control video length). + num_video_frames_per_chunk (`int`, *optional*): + Number of pixel frames generated per autoregressive chunk. + num_conditional_frames (`int`, *optional*, defaults to 1): + Number of frames each chunk reuses from the previous chunk's tail. + prompt (`str`): + The text prompt that guides Cosmos3 generation. + negative_prompt (`str`, *optional*): + The negative text prompt used for classifier-free guidance. + fps (`float`, *optional*, defaults to 24.0): + Frame rate of the generated video. + use_system_prompt (`bool`, *optional*, defaults to True): + Whether to prepend the Cosmos3 system prompt. + add_resolution_template (`bool`, *optional*, defaults to True): + Whether to add resolution metadata to the prompt. + add_duration_template (`bool`, *optional*, defaults to True): + Whether to add duration metadata to the prompt. + action (`CosmosActionCondition`, *optional*): + Action-conditioning metadata and its reference visual input. + video (`None`, *optional*): + Reference video for video-to-video conditioning. + condition_frame_indexes_vision (`tuple | list`, *optional*, defaults to (0, 1)): + Latent-frame indexes to preserve from the conditioning video. + condition_video_keep (`str`, *optional*, defaults to first): + Which end of a longer conditioning video to use: `first` or `last`. + image (`None`, *optional*): + Reference image for image-to-video conditioning. + chunk_id (`int`, *optional*, defaults to 0): + Index of the current chunk. + previous_output (`None`, *optional*): + Decoded pixels of the previous chunk, used to seed later chunks. + num_first_chunk_conditional_frames (`int`, *optional*, defaults to 0): + Number of frames the first chunk reuses from the input video. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + share_vision_temporal_positions (`bool`, *optional*, defaults to True): + Whether control and target items share vision temporal positions. + num_inference_steps (`int`): + The number of denoising steps. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + guidance_scale (`float`, *optional*, defaults to 6.0): + Scale for text classifier-free guidance. + control_guidance (`float`, *optional*, defaults to 1.0): + Scale for the control (structural) guidance axis. + guidance_interval (`tuple`, *optional*): + Timestep interval [lo, hi] over which text guidance is active (None = always). + control_guidance_interval (`tuple`, *optional*): + Timestep interval [lo, hi] over which control guidance is active (None = always). + output_chunks (`list`, *optional*): + Decoded pixel chunks accumulated so far. + x0_tokens_vision (`Tensor`, *optional*): + Vision latents encoded from the conditioning image or video. + vision_condition_frames (`list`, *optional*): + Latent-frame indexes fixed by visual conditioning. + latents (`Tensor`): + Pre-generated noisy vision latents. + sound_latents (`Tensor`, *optional*): + Pre-generated noisy sound latents. + action_condition_frame_indexes (`list`, *optional*): + Action-frame indexes fixed by action conditioning. + action_latents (`Tensor`, *optional*): + Pre-generated noisy action latents. + enable_sound (`bool`, *optional*, defaults to False): + Whether to generate a synchronized sound track. + output_type (`str`, *optional*, defaults to pil): + Output format: 'pil', 'np', 'pt'. + + Outputs: + videos (`list`): + The generated videos. + sound (`Tensor`): + Generated waveform. + sampling_rate (`int`): + Sample rate of the generated waveform in Hz. + action (`list`): + Generated action vectors. + """ + + model_name = "cosmos3-omni" + block_classes = [ + Cosmos3AutoTextEncoderStep, + Cosmos3AutoVaeEncoderStep, + Cosmos3AutoCoreDenoiseStep, + Cosmos3AutoDecodeStep, + Cosmos3ActionOutputStep, + ] + block_names = ["text_encoder", "vae_encoder", "denoise", "decode", "after_decode"] + _workflow_map = { + "text2image": {"prompt": True, "num_frames": 1}, + "text2video": {"prompt": True}, + "image2video": {"prompt": True, "image": True}, + "video2video": {"prompt": True, "video": True}, + "text2video_with_sound": {"prompt": True, "enable_sound": True}, + "image2video_with_sound": {"prompt": True, "image": True, "enable_sound": True}, + "video2video_with_sound": {"prompt": True, "video": True, "enable_sound": True}, + "action_policy": {"prompt": True, "action": True}, + "action_forward_dynamics": {"prompt": True, "action": True}, + "action_inverse_dynamics": {"prompt": True, "action": True}, + "transfer": {"prompt": True, "control_videos": True}, + } + + @property + def description(self): + return "Modular pipeline blocks for Cosmos3 generation modes." + + @property + def outputs(self): + return [ + OutputParam.template("videos"), + OutputParam("sound", type_hint=torch.Tensor, description="Generated waveform."), + OutputParam("sampling_rate", type_hint=int, description="Sample rate of the generated waveform in Hz."), + OutputParam("action", type_hint=list[torch.Tensor], description="Generated action vectors."), + ] diff --git a/src/diffusers/modular_pipelines/cosmos/modular_pipeline.py b/src/diffusers/modular_pipelines/cosmos/modular_pipeline.py new file mode 100644 index 000000000000..4628f6bec600 --- /dev/null +++ b/src/diffusers/modular_pipelines/cosmos/modular_pipeline.py @@ -0,0 +1,118 @@ +import torch + +from ...pipelines.cosmos.pipeline_cosmos3_omni import Cosmos3OmniPipeline, CosmosSafetyChecker +from ..modular_pipeline import ModularPipeline + + +class Cosmos3OmniModularPipeline(ModularPipeline): + """ + A ModularPipeline for Cosmos 3 omni generation. + """ + + default_blocks_name = "Cosmos3OmniBlocks" + + duration_template = "The video is {duration:.1f} seconds long and is of {fps:.0f} FPS." + image_resolution_template = "This image is of {height}x{width} resolution." + video_resolution_template = "This video is of {height}x{width} resolution." + inverse_duration_template = "The video is not {duration:.1f} seconds long and is not of {fps:.0f} FPS." + inverse_image_resolution_template = "This image is not of {height}x{width} resolution." + inverse_video_resolution_template = "This video is not of {height}x{width} resolution." + + @property + def vae_scale_factor_spatial(self): + if getattr(self, "vae", None) is not None: + return int(self.vae.config.scale_factor_spatial) + return 16 + + @property + def vae_scale_factor_temporal(self): + if getattr(self, "vae", None) is not None: + return int(self.vae.config.scale_factor_temporal) + return 4 + + @property + def num_channels_latents(self): + if getattr(self, "transformer", None) is not None: + return int(self.transformer.config.latent_channel) + return 48 + + @property + def sound_sampling_rate(self): + if getattr(self, "sound_tokenizer", None) is not None: + return int(self.sound_tokenizer.config.sampling_rate) + return 48000 + + @property + def sound_hop_size(self): + if getattr(self, "sound_tokenizer", None) is not None: + return int(self.sound_tokenizer._hop_size) + return 1920 + + @property + def _vae_latents_mean(self): + return torch.tensor(self.vae.config.latents_mean, dtype=self.vae.dtype) + + @property + def _vae_latents_inv_std(self): + return 1.0 / torch.tensor(self.vae.config.latents_std, dtype=self.vae.dtype) + + @property + def llm_special_tokens(self): + if getattr(self, "text_tokenizer", None) is None: + return None + return { + "start_of_generation": self.text_tokenizer.convert_tokens_to_ids("<|vision_start|>"), + "eos_token_id": self.text_tokenizer.eos_token_id, + } + + def enable_safety_checker(self, safety_checker=None): + if safety_checker is not None: + self.safety_checker = safety_checker + elif getattr(self, "safety_checker", None) is None: + self.safety_checker = CosmosSafetyChecker() + self._is_safety_checker_enabled = True + + def disable_safety_checker(self): + self._is_safety_checker_enabled = False + + @property + def requires_safety_checker(self): + return getattr(self, "_is_safety_checker_enabled", True) + + def _encode_video(self, x): + return Cosmos3OmniPipeline._encode_video(self, x) + + def decode_sound(self, latent): + return Cosmos3OmniPipeline.decode_sound(self, latent) + + def _prepare_text_segment(self, input_ids, device): + return Cosmos3OmniPipeline._prepare_text_segment(self, input_ids, device) + + def _prepare_vision_segment(self, *args, **kwargs): + return Cosmos3OmniPipeline._prepare_vision_segment(self, *args, **kwargs) + + def _prepare_sound_segment(self, *args, **kwargs): + return Cosmos3OmniPipeline._prepare_sound_segment(self, *args, **kwargs) + + def _prepare_action_segment(self, *args, **kwargs): + return Cosmos3OmniPipeline._prepare_action_segment(self, *args, **kwargs) + + def _prepare_action_video_conditioning(self, *args, **kwargs): + return Cosmos3OmniPipeline._prepare_action_video_conditioning(self, *args, **kwargs) + + def _remove_action_video_padding_from_latent(self, *args, **kwargs): + return Cosmos3OmniPipeline._remove_action_video_padding_from_latent(self, *args, **kwargs) + + @staticmethod + def _build_action_json_prompt(*args, **kwargs): + return Cosmos3OmniPipeline._build_action_json_prompt(*args, **kwargs) + + def tokenize_prompt(self, *args, **kwargs): + return Cosmos3OmniPipeline.tokenize_prompt(self, *args, **kwargs) + + @staticmethod + def _mask_velocity_predictions(*args, **kwargs): + return Cosmos3OmniPipeline._mask_velocity_predictions(*args, **kwargs) + + def _apply_video_safety_check(self, *args, **kwargs): + return Cosmos3OmniPipeline._apply_video_safety_check(self, *args, **kwargs) diff --git a/src/diffusers/modular_pipelines/modular_pipeline.py b/src/diffusers/modular_pipelines/modular_pipeline.py index 2d661028acf6..d43825860d8e 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/modular_pipeline.py @@ -133,6 +133,7 @@ def _helios_pyramid_map_fn(config_dict=None): ("qwenimage-layered", _create_default_map_fn("QwenImageLayeredModularPipeline")), ("anima", _create_default_map_fn("AnimaModularPipeline")), ("z-image", _create_default_map_fn("ZImageModularPipeline")), + ("cosmos3-omni", _create_default_map_fn("Cosmos3OmniModularPipeline")), ("helios", _create_default_map_fn("HeliosModularPipeline")), ("helios-pyramid", _helios_pyramid_map_fn), ("hunyuan-video-1.5", _create_default_map_fn("HunyuanVideo15ModularPipeline")), diff --git a/src/diffusers/pipelines/__init__.py b/src/diffusers/pipelines/__init__.py index bcd48f9649c5..7b0c43727975 100644 --- a/src/diffusers/pipelines/__init__.py +++ b/src/diffusers/pipelines/__init__.py @@ -346,7 +346,12 @@ "LTX2ImageToVideoPipeline", "LTX2LatentUpsamplePipeline", ] - _import_structure["joyimage"] = ["JoyImageEditPipeline", "JoyImageEditPipelineOutput"] + _import_structure["joyimage"] = [ + "JoyImageEditPipeline", + "JoyImageEditPipelineOutput", + "JoyImageEditPlusPipeline", + "JoyImageEditPlusPipelineOutput", + ] _import_structure["lumina"] = ["LuminaPipeline", "LuminaText2ImgPipeline"] _import_structure["lumina2"] = ["Lumina2Pipeline", "Lumina2Text2ImgPipeline"] _import_structure["lucy"] = ["LucyEditPipeline"] @@ -760,7 +765,12 @@ from .hunyuan_video1_5 import HunyuanVideo15ImageToVideoPipeline, HunyuanVideo15Pipeline from .hunyuandit import HunyuanDiTPipeline from .ideogram4 import Ideogram4Pipeline, Ideogram4PromptEnhancerHead - from .joyimage import JoyImageEditPipeline, JoyImageEditPipelineOutput + from .joyimage import ( + JoyImageEditPipeline, + JoyImageEditPipelineOutput, + JoyImageEditPlusPipeline, + JoyImageEditPlusPipelineOutput, + ) from .kandinsky import ( KandinskyCombinedPipeline, KandinskyImg2ImgCombinedPipeline, diff --git a/src/diffusers/pipelines/allegro/pipeline_allegro.py b/src/diffusers/pipelines/allegro/pipeline_allegro.py index 5949ed407661..9d2d2aa8bd09 100644 --- a/src/diffusers/pipelines/allegro/pipeline_allegro.py +++ b/src/diffusers/pipelines/allegro/pipeline_allegro.py @@ -68,7 +68,7 @@ >>> vae = AutoencoderKLAllegro.from_pretrained("rhymes-ai/Allegro", subfolder="vae", torch_dtype=torch.float32) >>> pipe = AllegroPipeline.from_pretrained("rhymes-ai/Allegro", vae=vae, torch_dtype=torch.bfloat16).to("cuda") - >>> pipe.enable_vae_tiling() + >>> pipe.vae.enable_tiling() >>> prompt = ( ... "A seaside harbor with bright sunlight and sparkling seawater, with many boats in the water. From an aerial view, " @@ -646,59 +646,6 @@ def _prepare_rotary_positional_embeddings( return (freqs_t, freqs_h, freqs_w), (grid_t, grid_h, grid_w) - def enable_vae_slicing(self): - r""" - Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to - compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. - """ - depr_message = f"Calling `enable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_slicing()`." - deprecate( - "enable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.enable_slicing() - - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_slicing()`." - deprecate( - "disable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.disable_slicing() - - def enable_vae_tiling(self): - r""" - Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to - compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow - processing larger images. - """ - depr_message = f"Calling `enable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_tiling()`." - deprecate( - "enable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.enable_tiling() - - def disable_vae_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_tiling()`." - deprecate( - "disable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.disable_tiling() - @property def guidance_scale(self): return self._guidance_scale diff --git a/src/diffusers/pipelines/animatediff/pipeline_animatediff_sdxl.py b/src/diffusers/pipelines/animatediff/pipeline_animatediff_sdxl.py index efcb64b96ec8..a221f3e1efeb 100644 --- a/src/diffusers/pipelines/animatediff/pipeline_animatediff_sdxl.py +++ b/src/diffusers/pipelines/animatediff/pipeline_animatediff_sdxl.py @@ -101,8 +101,8 @@ ... ).to("cuda") >>> # enable memory savings - >>> pipe.enable_vae_slicing() - >>> pipe.enable_vae_tiling() + >>> pipe.vae.enable_slicing() + >>> pipe.vae.enable_tiling() >>> output = pipe( ... prompt="a panda surfing in the ocean, realistic, high quality", diff --git a/src/diffusers/pipelines/audioldm2/pipeline_audioldm2.py b/src/diffusers/pipelines/audioldm2/pipeline_audioldm2.py index 6fb02433dace..110ddd1bfef3 100644 --- a/src/diffusers/pipelines/audioldm2/pipeline_audioldm2.py +++ b/src/diffusers/pipelines/audioldm2/pipeline_audioldm2.py @@ -34,7 +34,6 @@ from ...models import AutoencoderKL from ...schedulers import KarrasDiffusionSchedulers from ...utils import ( - deprecate, is_accelerate_available, is_accelerate_version, is_librosa_available, @@ -223,34 +222,6 @@ def __init__( ) self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8 - # Copied from diffusers.pipelines.pipeline_utils.StableDiffusionMixin.enable_vae_slicing - def enable_vae_slicing(self): - r""" - Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to - compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. - """ - depr_message = f"Calling `enable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_slicing()`." - deprecate( - "enable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.enable_slicing() - - # Copied from diffusers.pipelines.pipeline_utils.StableDiffusionMixin.disable_vae_slicing - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_slicing()`." - deprecate( - "disable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.disable_slicing() - def enable_model_cpu_offload(self, gpu_id: int | None = None, device: torch.device | str = "cuda"): r""" Offloads all models to CPU using accelerate, reducing memory usage with a low impact on performance. Compared diff --git a/src/diffusers/pipelines/auto_pipeline.py b/src/diffusers/pipelines/auto_pipeline.py index 15c23039c1cb..5fe97a969a3a 100644 --- a/src/diffusers/pipelines/auto_pipeline.py +++ b/src/diffusers/pipelines/auto_pipeline.py @@ -41,6 +41,7 @@ StableDiffusion3ControlNetInpaintingPipeline, StableDiffusion3ControlNetPipeline, ) +from .cosmos import Cosmos3OmniPipeline from .deepfloyd_if import IFImg2ImgPipeline, IFInpaintingPipeline, IFPipeline from .deprecated.wuerstchen import WuerstchenCombinedPipeline, WuerstchenDecoderPipeline from .flux import ( @@ -270,6 +271,12 @@ ] ) +AUTO_CONDITION2VIDEO_PIPELINES_MAPPING = OrderedDict( + [ + ("cosmos3-omni", Cosmos3OmniPipeline), + ] +) + AUTO_IMAGE2VIDEO_PIPELINES_MAPPING = OrderedDict( [ ("anyflow-far", AnyFlowFARPipeline), @@ -318,6 +325,7 @@ AUTO_IMAGE2IMAGE_PIPELINES_MAPPING, AUTO_INPAINT_PIPELINES_MAPPING, AUTO_TEXT2VIDEO_PIPELINES_MAPPING, + AUTO_CONDITION2VIDEO_PIPELINES_MAPPING, AUTO_IMAGE2VIDEO_PIPELINES_MAPPING, AUTO_VIDEO2VIDEO_PIPELINES_MAPPING, AUTO_TEXT2AUDIO_PIPELINES_MAPPING, diff --git a/src/diffusers/pipelines/chroma/pipeline_chroma.py b/src/diffusers/pipelines/chroma/pipeline_chroma.py index bc782107022d..3d91f43b07cf 100644 --- a/src/diffusers/pipelines/chroma/pipeline_chroma.py +++ b/src/diffusers/pipelines/chroma/pipeline_chroma.py @@ -25,7 +25,6 @@ from ...schedulers import FlowMatchEulerDiscreteScheduler from ...utils import ( USE_PEFT_BACKEND, - deprecate, is_torch_xla_available, logging, replace_example_docstring, @@ -507,59 +506,6 @@ def _unpack_latents(latents, height, width, vae_scale_factor): return latents - def enable_vae_slicing(self): - r""" - Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to - compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. - """ - depr_message = f"Calling `enable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_slicing()`." - deprecate( - "enable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.enable_slicing() - - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_slicing()`." - deprecate( - "disable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.disable_slicing() - - def enable_vae_tiling(self): - r""" - Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to - compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow - processing larger images. - """ - depr_message = f"Calling `enable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_tiling()`." - deprecate( - "enable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.enable_tiling() - - def disable_vae_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_tiling()`." - deprecate( - "disable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.disable_tiling() - # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.prepare_latents def prepare_latents( self, diff --git a/src/diffusers/pipelines/chroma/pipeline_chroma_img2img.py b/src/diffusers/pipelines/chroma/pipeline_chroma_img2img.py index 6dad6a481c5a..4a2d4cca049e 100644 --- a/src/diffusers/pipelines/chroma/pipeline_chroma_img2img.py +++ b/src/diffusers/pipelines/chroma/pipeline_chroma_img2img.py @@ -25,7 +25,6 @@ from ...schedulers import FlowMatchEulerDiscreteScheduler from ...utils import ( USE_PEFT_BACKEND, - deprecate, is_torch_xla_available, logging, replace_example_docstring, @@ -539,59 +538,6 @@ def _unpack_latents(latents, height, width, vae_scale_factor): return latents - def enable_vae_slicing(self): - r""" - Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to - compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. - """ - depr_message = f"Calling `enable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_slicing()`." - deprecate( - "enable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.enable_slicing() - - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_slicing()`." - deprecate( - "disable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.disable_slicing() - - def enable_vae_tiling(self): - r""" - Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to - compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow - processing larger images. - """ - depr_message = f"Calling `enable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_tiling()`." - deprecate( - "enable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.enable_tiling() - - def disable_vae_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_tiling()`." - deprecate( - "disable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.disable_tiling() - # Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3_img2img.StableDiffusion3Img2ImgPipeline.get_timesteps def get_timesteps(self, num_inference_steps, strength, device): # get the original timestep using init_timestep diff --git a/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py b/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py index 7381aafc1777..12405755309f 100644 --- a/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py +++ b/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py @@ -398,6 +398,9 @@ def __init__( # Image preprocessor for caller-supplied conditioning frames (PIL / tensor / numpy). self.vae_scale_factor_spatial = int(self.vae.config.scale_factor_spatial) if getattr(self, "vae", None) else 16 + self.vae_scale_factor_temporal = ( + int(self.vae.config.scale_factor_temporal) if getattr(self, "vae", None) else 4 + ) self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial, resample="bilinear") self.llm_special_tokens = { @@ -545,7 +548,7 @@ def _prepare_vision_segment( reset_spatial_indices=config.unified_3d_mrope_reset_spatial_ids, fps=effective_fps, base_fps=float(config.base_fps), - temporal_compression_factor=self.vae.config.scale_factor_temporal, + temporal_compression_factor=self.vae_scale_factor_temporal, ) return { @@ -628,7 +631,7 @@ def _prepare_action_segment( fps=effective_fps, base_fps=float(config.base_fps), temporal_compression_factor=1, - base_temporal_compression_factor=self.vae.config.scale_factor_temporal, + base_temporal_compression_factor=self.vae_scale_factor_temporal, start_frame_offset=1, ) diff --git a/src/diffusers/pipelines/deprecated/unidiffuser/pipeline_unidiffuser.py b/src/diffusers/pipelines/deprecated/unidiffuser/pipeline_unidiffuser.py index 7e55075cc209..3b3db863744b 100644 --- a/src/diffusers/pipelines/deprecated/unidiffuser/pipeline_unidiffuser.py +++ b/src/diffusers/pipelines/deprecated/unidiffuser/pipeline_unidiffuser.py @@ -226,63 +226,6 @@ def _infer_mode(self, prompt, prompt_embeds, image, latents, prompt_latents, vae return mode - # Copied from diffusers.pipelines.pipeline_utils.StableDiffusionMixin.enable_vae_slicing - def enable_vae_slicing(self): - r""" - Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to - compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. - """ - depr_message = f"Calling `enable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_slicing()`." - deprecate( - "enable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.enable_slicing() - - # Copied from diffusers.pipelines.pipeline_utils.StableDiffusionMixin.disable_vae_slicing - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_slicing()`." - deprecate( - "disable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.disable_slicing() - - # Copied from diffusers.pipelines.pipeline_utils.StableDiffusionMixin.enable_vae_tiling - def enable_vae_tiling(self): - r""" - Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to - compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow - processing larger images. - """ - depr_message = f"Calling `enable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_tiling()`." - deprecate( - "enable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.enable_tiling() - - # Copied from diffusers.pipelines.pipeline_utils.StableDiffusionMixin.disable_vae_tiling - def disable_vae_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_tiling()`." - deprecate( - "disable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.disable_tiling() - # Functions to manually set the mode def set_text_mode(self): r"""Manually set the generation mode to unconditional ("marginal") text generation.""" diff --git a/src/diffusers/pipelines/flux/pipeline_flux.py b/src/diffusers/pipelines/flux/pipeline_flux.py index 34cbf0faa667..d3e0682c5419 100644 --- a/src/diffusers/pipelines/flux/pipeline_flux.py +++ b/src/diffusers/pipelines/flux/pipeline_flux.py @@ -32,7 +32,6 @@ from ...schedulers import FlowMatchEulerDiscreteScheduler from ...utils import ( USE_PEFT_BACKEND, - deprecate, is_torch_xla_available, logging, replace_example_docstring, @@ -541,59 +540,6 @@ def _unpack_latents(latents, height, width, vae_scale_factor): return latents - def enable_vae_slicing(self): - r""" - Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to - compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. - """ - depr_message = f"Calling `enable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_slicing()`." - deprecate( - "enable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.enable_slicing() - - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_slicing()`." - deprecate( - "disable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.disable_slicing() - - def enable_vae_tiling(self): - r""" - Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to - compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow - processing larger images. - """ - depr_message = f"Calling `enable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_tiling()`." - deprecate( - "enable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.enable_tiling() - - def disable_vae_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_tiling()`." - deprecate( - "disable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.disable_tiling() - def prepare_latents( self, batch_size, diff --git a/src/diffusers/pipelines/flux/pipeline_flux_control.py b/src/diffusers/pipelines/flux/pipeline_flux_control.py index 84e21b59ce58..46671c44cca8 100644 --- a/src/diffusers/pipelines/flux/pipeline_flux_control.py +++ b/src/diffusers/pipelines/flux/pipeline_flux_control.py @@ -26,7 +26,6 @@ from ...schedulers import FlowMatchEulerDiscreteScheduler from ...utils import ( USE_PEFT_BACKEND, - deprecate, is_torch_xla_available, logging, replace_example_docstring, @@ -492,59 +491,6 @@ def _unpack_latents(latents, height, width, vae_scale_factor): return latents - def enable_vae_slicing(self): - r""" - Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to - compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. - """ - depr_message = f"Calling `enable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_slicing()`." - deprecate( - "enable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.enable_slicing() - - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_slicing()`." - deprecate( - "disable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.disable_slicing() - - def enable_vae_tiling(self): - r""" - Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to - compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow - processing larger images. - """ - depr_message = f"Calling `enable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_tiling()`." - deprecate( - "enable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.enable_tiling() - - def disable_vae_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_tiling()`." - deprecate( - "disable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.disable_tiling() - # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.prepare_latents def prepare_latents( self, diff --git a/src/diffusers/pipelines/flux/pipeline_flux_control_inpaint.py b/src/diffusers/pipelines/flux/pipeline_flux_control_inpaint.py index cd4ee9fe7611..15e27653c3e2 100644 --- a/src/diffusers/pipelines/flux/pipeline_flux_control_inpaint.py +++ b/src/diffusers/pipelines/flux/pipeline_flux_control_inpaint.py @@ -35,7 +35,6 @@ from ...schedulers import FlowMatchEulerDiscreteScheduler from ...utils import ( USE_PEFT_BACKEND, - deprecate, is_torch_xla_available, logging, replace_example_docstring, @@ -573,59 +572,6 @@ def _unpack_latents(latents, height, width, vae_scale_factor): return latents - def enable_vae_slicing(self): - r""" - Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to - compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. - """ - depr_message = f"Calling `enable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_slicing()`." - deprecate( - "enable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.enable_slicing() - - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_slicing()`." - deprecate( - "disable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.disable_slicing() - - def enable_vae_tiling(self): - r""" - Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to - compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow - processing larger images. - """ - depr_message = f"Calling `enable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_tiling()`." - deprecate( - "enable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.enable_tiling() - - def disable_vae_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_tiling()`." - deprecate( - "disable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.disable_tiling() - def prepare_latents( self, image, diff --git a/src/diffusers/pipelines/flux/pipeline_flux_fill.py b/src/diffusers/pipelines/flux/pipeline_flux_fill.py index 4098213cc894..ab4431b5b768 100644 --- a/src/diffusers/pipelines/flux/pipeline_flux_fill.py +++ b/src/diffusers/pipelines/flux/pipeline_flux_fill.py @@ -26,7 +26,6 @@ from ...schedulers import FlowMatchEulerDiscreteScheduler from ...utils import ( USE_PEFT_BACKEND, - deprecate, is_torch_xla_available, logging, replace_example_docstring, @@ -629,59 +628,6 @@ def _unpack_latents(latents, height, width, vae_scale_factor): return latents - def enable_vae_slicing(self): - r""" - Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to - compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. - """ - depr_message = f"Calling `enable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_slicing()`." - deprecate( - "enable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.enable_slicing() - - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_slicing()`." - deprecate( - "disable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.disable_slicing() - - def enable_vae_tiling(self): - r""" - Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to - compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow - processing larger images. - """ - depr_message = f"Calling `enable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_tiling()`." - deprecate( - "enable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.enable_tiling() - - def disable_vae_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_tiling()`." - deprecate( - "disable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.disable_tiling() - # Copied from diffusers.pipelines.flux.pipeline_flux_img2img.FluxImg2ImgPipeline.prepare_latents def prepare_latents( self, diff --git a/src/diffusers/pipelines/flux/pipeline_flux_img2img.py b/src/diffusers/pipelines/flux/pipeline_flux_img2img.py index 18206ec36600..94582a84cf84 100644 --- a/src/diffusers/pipelines/flux/pipeline_flux_img2img.py +++ b/src/diffusers/pipelines/flux/pipeline_flux_img2img.py @@ -33,7 +33,6 @@ from ...schedulers import FlowMatchEulerDiscreteScheduler from ...utils import ( USE_PEFT_BACKEND, - deprecate, is_torch_xla_available, logging, replace_example_docstring, @@ -608,63 +607,6 @@ def _unpack_latents(latents, height, width, vae_scale_factor): return latents - # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.enable_vae_slicing - def enable_vae_slicing(self): - r""" - Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to - compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. - """ - depr_message = f"Calling `enable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_slicing()`." - deprecate( - "enable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.enable_slicing() - - # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.disable_vae_slicing - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_slicing()`." - deprecate( - "disable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.disable_slicing() - - # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.enable_vae_tiling - def enable_vae_tiling(self): - r""" - Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to - compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow - processing larger images. - """ - depr_message = f"Calling `enable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_tiling()`." - deprecate( - "enable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.enable_tiling() - - # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.disable_vae_tiling - def disable_vae_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_tiling()`." - deprecate( - "disable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.disable_tiling() - def prepare_latents( self, image, diff --git a/src/diffusers/pipelines/flux/pipeline_flux_kontext.py b/src/diffusers/pipelines/flux/pipeline_flux_kontext.py index e32bfecfcdad..849d9686de62 100644 --- a/src/diffusers/pipelines/flux/pipeline_flux_kontext.py +++ b/src/diffusers/pipelines/flux/pipeline_flux_kontext.py @@ -32,7 +32,6 @@ from ...schedulers import FlowMatchEulerDiscreteScheduler from ...utils import ( USE_PEFT_BACKEND, - deprecate, is_torch_xla_available, logging, replace_example_docstring, @@ -609,63 +608,6 @@ def _encode_vae_image(self, image: torch.Tensor, generator: torch.Generator): return image_latents - # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.enable_vae_slicing - def enable_vae_slicing(self): - r""" - Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to - compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. - """ - depr_message = f"Calling `enable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_slicing()`." - deprecate( - "enable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.enable_slicing() - - # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.disable_vae_slicing - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_slicing()`." - deprecate( - "disable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.disable_slicing() - - # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.enable_vae_tiling - def enable_vae_tiling(self): - r""" - Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to - compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow - processing larger images. - """ - depr_message = f"Calling `enable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_tiling()`." - deprecate( - "enable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.enable_tiling() - - # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.disable_vae_tiling - def disable_vae_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_tiling()`." - deprecate( - "disable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.disable_tiling() - def prepare_latents( self, image: torch.Tensor | None, diff --git a/src/diffusers/pipelines/flux/pipeline_flux_kontext_inpaint.py b/src/diffusers/pipelines/flux/pipeline_flux_kontext_inpaint.py index c85299eedcd3..9a79aafd8760 100644 --- a/src/diffusers/pipelines/flux/pipeline_flux_kontext_inpaint.py +++ b/src/diffusers/pipelines/flux/pipeline_flux_kontext_inpaint.py @@ -22,7 +22,6 @@ from ...schedulers import FlowMatchEulerDiscreteScheduler from ...utils import ( USE_PEFT_BACKEND, - deprecate, is_torch_xla_available, logging, replace_example_docstring, @@ -683,63 +682,6 @@ def _encode_vae_image(self, image: torch.Tensor, generator: torch.Generator): return image_latents - # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.enable_vae_slicing - def enable_vae_slicing(self): - r""" - Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to - compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. - """ - depr_message = f"Calling `enable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_slicing()`." - deprecate( - "enable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.enable_slicing() - - # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.disable_vae_slicing - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_slicing()`." - deprecate( - "disable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.disable_slicing() - - # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.enable_vae_tiling - def enable_vae_tiling(self): - r""" - Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to - compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow - processing larger images. - """ - depr_message = f"Calling `enable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_tiling()`." - deprecate( - "enable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.enable_tiling() - - # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.disable_vae_tiling - def disable_vae_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_tiling()`." - deprecate( - "disable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.disable_tiling() - def prepare_latents( self, image: torch.Tensor | None, diff --git a/src/diffusers/pipelines/hidream_image/pipeline_hidream_image.py b/src/diffusers/pipelines/hidream_image/pipeline_hidream_image.py index 1c73dfacccdb..d723c7941ae9 100644 --- a/src/diffusers/pipelines/hidream_image/pipeline_hidream_image.py +++ b/src/diffusers/pipelines/hidream_image/pipeline_hidream_image.py @@ -516,59 +516,6 @@ def encode_prompt( negative_pooled_prompt_embeds, ) - def enable_vae_slicing(self): - r""" - Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to - compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. - """ - depr_message = f"Calling `enable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_slicing()`." - deprecate( - "enable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.enable_slicing() - - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_slicing()`." - deprecate( - "disable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.disable_slicing() - - def enable_vae_tiling(self): - r""" - Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to - compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow - processing larger images. - """ - depr_message = f"Calling `enable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_tiling()`." - deprecate( - "enable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.enable_tiling() - - def disable_vae_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_tiling()`." - deprecate( - "disable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.disable_tiling() - def check_inputs( self, prompt, diff --git a/src/diffusers/pipelines/hunyuan_video/pipeline_hunyuan_skyreels_image2video.py b/src/diffusers/pipelines/hunyuan_video/pipeline_hunyuan_skyreels_image2video.py index b5b4ff9bcd85..bd54f2563b52 100644 --- a/src/diffusers/pipelines/hunyuan_video/pipeline_hunyuan_skyreels_image2video.py +++ b/src/diffusers/pipelines/hunyuan_video/pipeline_hunyuan_skyreels_image2video.py @@ -24,7 +24,7 @@ from ...loaders import HunyuanVideoLoraLoaderMixin from ...models import AutoencoderKLHunyuanVideo, HunyuanVideoTransformer3DModel from ...schedulers import FlowMatchEulerDiscreteScheduler -from ...utils import deprecate, is_torch_xla_available, logging, replace_example_docstring +from ...utils import is_torch_xla_available, logging, replace_example_docstring from ...utils.torch_utils import randn_tensor from ...video_processor import VideoProcessor from ..pipeline_utils import DiffusionPipeline @@ -458,59 +458,6 @@ def prepare_latents( return latents, image_latents - def enable_vae_slicing(self): - r""" - Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to - compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. - """ - depr_message = f"Calling `enable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_slicing()`." - deprecate( - "enable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.enable_slicing() - - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_slicing()`." - deprecate( - "disable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.disable_slicing() - - def enable_vae_tiling(self): - r""" - Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to - compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow - processing larger images. - """ - depr_message = f"Calling `enable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_tiling()`." - deprecate( - "enable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.enable_tiling() - - def disable_vae_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_tiling()`." - deprecate( - "disable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.disable_tiling() - @property def guidance_scale(self): return self._guidance_scale diff --git a/src/diffusers/pipelines/hunyuan_video/pipeline_hunyuan_video.py b/src/diffusers/pipelines/hunyuan_video/pipeline_hunyuan_video.py index 5b8cff2ca0c5..5e785c38891c 100644 --- a/src/diffusers/pipelines/hunyuan_video/pipeline_hunyuan_video.py +++ b/src/diffusers/pipelines/hunyuan_video/pipeline_hunyuan_video.py @@ -23,7 +23,7 @@ from ...loaders import HunyuanVideoLoraLoaderMixin from ...models import AutoencoderKLHunyuanVideo, HunyuanVideoTransformer3DModel from ...schedulers import FlowMatchEulerDiscreteScheduler -from ...utils import deprecate, is_torch_xla_available, logging, replace_example_docstring +from ...utils import is_torch_xla_available, logging, replace_example_docstring from ...utils.torch_utils import randn_tensor from ...video_processor import VideoProcessor from ..pipeline_utils import DiffusionPipeline @@ -415,59 +415,6 @@ def prepare_latents( latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) return latents - def enable_vae_slicing(self): - r""" - Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to - compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. - """ - depr_message = f"Calling `enable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_slicing()`." - deprecate( - "enable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.enable_slicing() - - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_slicing()`." - deprecate( - "disable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.disable_slicing() - - def enable_vae_tiling(self): - r""" - Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to - compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow - processing larger images. - """ - depr_message = f"Calling `enable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_tiling()`." - deprecate( - "enable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.enable_tiling() - - def disable_vae_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_tiling()`." - deprecate( - "disable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.disable_tiling() - @property def guidance_scale(self): return self._guidance_scale diff --git a/src/diffusers/pipelines/hunyuan_video/pipeline_hunyuan_video_framepack.py b/src/diffusers/pipelines/hunyuan_video/pipeline_hunyuan_video_framepack.py index 515b530d1037..349481492ac0 100644 --- a/src/diffusers/pipelines/hunyuan_video/pipeline_hunyuan_video_framepack.py +++ b/src/diffusers/pipelines/hunyuan_video/pipeline_hunyuan_video_framepack.py @@ -33,7 +33,7 @@ from ...loaders import HunyuanVideoLoraLoaderMixin from ...models import AutoencoderKLHunyuanVideo, HunyuanVideoFramepackTransformer3DModel from ...schedulers import FlowMatchEulerDiscreteScheduler -from ...utils import deprecate, is_torch_xla_available, logging, replace_example_docstring +from ...utils import is_torch_xla_available, logging, replace_example_docstring from ...utils.torch_utils import randn_tensor from ...video_processor import VideoProcessor from ..pipeline_utils import DiffusionPipeline @@ -563,59 +563,6 @@ def prepare_image_latents( latents = latents * self.vae.config.scaling_factor return latents.to(device=device, dtype=dtype) - def enable_vae_slicing(self): - r""" - Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to - compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. - """ - depr_message = f"Calling `enable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_slicing()`." - deprecate( - "enable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.enable_slicing() - - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_slicing()`." - deprecate( - "disable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.disable_slicing() - - def enable_vae_tiling(self): - r""" - Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to - compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow - processing larger images. - """ - depr_message = f"Calling `enable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_tiling()`." - deprecate( - "enable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.enable_tiling() - - def disable_vae_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_tiling()`." - deprecate( - "disable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.disable_tiling() - @property def guidance_scale(self): return self._guidance_scale diff --git a/src/diffusers/pipelines/hunyuan_video/pipeline_hunyuan_video_image2video.py b/src/diffusers/pipelines/hunyuan_video/pipeline_hunyuan_video_image2video.py index 1c68be879013..13eb35386001 100644 --- a/src/diffusers/pipelines/hunyuan_video/pipeline_hunyuan_video_image2video.py +++ b/src/diffusers/pipelines/hunyuan_video/pipeline_hunyuan_video_image2video.py @@ -30,7 +30,7 @@ from ...loaders import HunyuanVideoLoraLoaderMixin from ...models import AutoencoderKLHunyuanVideo, HunyuanVideoTransformer3DModel from ...schedulers import FlowMatchEulerDiscreteScheduler -from ...utils import deprecate, is_torch_xla_available, logging, replace_example_docstring +from ...utils import is_torch_xla_available, logging, replace_example_docstring from ...utils.torch_utils import randn_tensor from ...video_processor import VideoProcessor from ..pipeline_utils import DiffusionPipeline @@ -598,59 +598,6 @@ def prepare_latents( return latents, image_latents - def enable_vae_slicing(self): - r""" - Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to - compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. - """ - depr_message = f"Calling `enable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_slicing()`." - deprecate( - "enable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.enable_slicing() - - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_slicing()`." - deprecate( - "disable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.disable_slicing() - - def enable_vae_tiling(self): - r""" - Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to - compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow - processing larger images. - """ - depr_message = f"Calling `enable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_tiling()`." - deprecate( - "enable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.enable_tiling() - - def disable_vae_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_tiling()`." - deprecate( - "disable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.disable_tiling() - @property def guidance_scale(self): return self._guidance_scale diff --git a/src/diffusers/pipelines/joyimage/__init__.py b/src/diffusers/pipelines/joyimage/__init__.py index 85b9246b22a6..a5faea9d9763 100644 --- a/src/diffusers/pipelines/joyimage/__init__.py +++ b/src/diffusers/pipelines/joyimage/__init__.py @@ -22,8 +22,8 @@ _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects)) else: _import_structure["pipeline_joyimage_edit"] = ["JoyImageEditPipeline"] - - _import_structure["pipeline_output"] = ["JoyImageEditPipelineOutput"] + _import_structure["pipeline_joyimage_edit_plus"] = ["JoyImageEditPlusPipeline"] + _import_structure["pipeline_output"] = ["JoyImageEditPipelineOutput", "JoyImageEditPlusPipelineOutput"] if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: try: @@ -34,7 +34,8 @@ from ...utils.dummy_torch_and_transformers_objects import * else: from .pipeline_joyimage_edit import JoyImageEditPipeline - from .pipeline_output import JoyImageEditPipelineOutput + from .pipeline_joyimage_edit_plus import JoyImageEditPlusPipeline + from .pipeline_output import JoyImageEditPipelineOutput, JoyImageEditPlusPipelineOutput else: import sys diff --git a/src/diffusers/pipelines/joyimage/pipeline_joyimage_edit_plus.py b/src/diffusers/pipelines/joyimage/pipeline_joyimage_edit_plus.py new file mode 100644 index 000000000000..ac8e01278e3e --- /dev/null +++ b/src/diffusers/pipelines/joyimage/pipeline_joyimage_edit_plus.py @@ -0,0 +1,770 @@ +# Copyright 2025 The JoyImage Team and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import inspect +from typing import Callable + +import torch +from PIL import Image +from transformers import ( + Qwen2Tokenizer, + Qwen3VLForConditionalGeneration, + Qwen3VLProcessor, +) + +from ...callbacks import MultiPipelineCallbacks, PipelineCallback +from ...models import AutoencoderKLWan +from ...models.transformers.transformer_joyimage_edit_plus import JoyImageEditPlusTransformer3DModel +from ...schedulers import FlowMatchEulerDiscreteScheduler +from ...utils import logging, replace_example_docstring +from ...utils.torch_utils import randn_tensor +from ..pipeline_utils import DiffusionPipeline +from .image_processor import JoyImageEditImageProcessor +from .pipeline_output import JoyImageEditPlusPipelineOutput + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +EXAMPLE_DOC_STRING = """ +Examples: + ```python + >>> import torch + >>> from diffusers import JoyImageEditPlusPipeline + >>> from diffusers.utils import load_image + + >>> model_id = "jdopensource/JoyAI-Image-Edit-Plus-Diffusers" + >>> pipe = JoyImageEditPlusPipeline.from_pretrained(model_id, torch_dtype=torch.bfloat16) + >>> pipe.to("cuda") + + >>> images = [ + ... load_image("dog.png"), + ... load_image("person.png"), + ... ] + >>> output = pipe( + ... images=images, + ... prompt="Let the person lovingly play with the dog.", + ... height=1024, + ... width=1024, + ... num_inference_steps=30, + ... guidance_scale=4.0, + ... generator=torch.manual_seed(42), + ... ) + >>> output.images[0].save("output.png") + ``` +""" + + +# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps +def retrieve_timesteps( + scheduler, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, + **kwargs, +): + r""" + Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles + custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`. + + Args: + scheduler (`SchedulerMixin`): + The scheduler to get timesteps from. + num_inference_steps (`int`): + The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps` + must be `None`. + device (`str` or `torch.device`, *optional*): + The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. + timesteps (`list[int]`, *optional*): + Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, + `num_inference_steps` and `sigmas` must be `None`. + sigmas (`list[float]`, *optional*): + Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, + `num_inference_steps` and `timesteps` must be `None`. + + Returns: + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + second element is the number of inference steps. + """ + if timesteps is not None and sigmas is not None: + raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values") + if timesteps is not None: + accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accepts_timesteps: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" timestep schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + elif sigmas is not None: + accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accept_sigmas: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" sigmas schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + else: + scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) + timesteps = scheduler.timesteps + return timesteps, num_inference_steps + + +class JoyImageEditPlusPipeline(DiffusionPipeline): + r""" + Diffusion pipeline for multi-image instruction-guided editing using JoyImage Edit Plus. + + Supports multiple reference images with different resolutions. Each reference image is independently VAE-encoded + and patchified, then concatenated with the target noise patches for joint denoising. + + Args: + scheduler ([`FlowMatchEulerDiscreteScheduler`]): + A scheduler to be used in combination with `transformer` to denoise the encoded image latents. + vae ([`AutoencoderKLWan`]): + Variational Auto-Encoder (VAE) model to encode and decode images to and from latent representations. + text_encoder ([`Qwen3VLForConditionalGeneration`]): + Multimodal text encoder for prompt encoding with inline image understanding. + tokenizer ([`Qwen2Tokenizer`]): + Tokenizer for text processing. + transformer ([`JoyImageEditPlusTransformer3DModel`]): + Conditional Transformer (MMDiT) architecture to denoise the encoded image latents. + processor ([`Qwen3VLProcessor`]): + Processor for multimodal inputs (text + images). + text_token_max_length (`int`, defaults to `2048`): + Maximum token length for text encoding. + """ + + model_cpu_offload_seq = "text_encoder->transformer->vae" + _callback_tensor_inputs = ["latents", "prompt_embeds"] + + def __init__( + self, + scheduler: FlowMatchEulerDiscreteScheduler, + vae: AutoencoderKLWan, + text_encoder: Qwen3VLForConditionalGeneration, + tokenizer: Qwen2Tokenizer, + transformer: JoyImageEditPlusTransformer3DModel, + processor: Qwen3VLProcessor, + text_token_max_length: int = 2048, + ): + super().__init__() + self.register_modules( + vae=vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + transformer=transformer, + scheduler=scheduler, + processor=processor, + ) + + self.text_token_max_length = text_token_max_length + + self.vae_scale_factor_temporal = self.vae.config.scale_factor_temporal if getattr(self, "vae", None) else 4 + self.vae_scale_factor_spatial = self.vae.config.scale_factor_spatial if getattr(self, "vae", None) else 8 + self.image_processor = JoyImageEditImageProcessor(vae_scale_factor=self.vae_scale_factor_spatial) + + self.prompt_template_encode = { + "multiple_images": ( + "<|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" + ), + } + self.prompt_template_encode_start_idx = { + "multiple_images": 34, + } + + # ------------------------------------------------------------------ + # Internal helpers + # ------------------------------------------------------------------ + + def _get_last_decoder_hidden_states(self, forward_fn, **kwargs): + """ + Run ``forward_fn(**kwargs)`` while capturing the **pre-norm** output of the last decoder layer via a forward + hook. + + This model was trained on transformers 4.57, where ``Qwen3VLForConditionalGeneration``'s + ``@check_model_inputs`` decorator monkey-patched each decoder layer to collect ``hidden_states``. Because + ``Qwen3VLCausalLMOutputWithPast`` has no ``last_hidden_state`` field, ``tie_last_hidden_states`` had no effect + and ``hidden_states[-1]`` was the **pre-norm** output of the last decoder layer. + + Starting from https://github.com/huggingface/transformers/pull/42609 the CausalLM forward explicitly returns + ``hidden_states=outputs.hidden_states`` from the inner model. Combined with the subsequent + ``@check_model_inputs`` → ``@capture_outputs`` migration (transformers 5.x), ``hidden_states`` is now captured + at the ``Qwen3VLTextModel`` level where ``tie_last_hidden_states=True`` replaces ``hidden_states[-1]`` with the + **post-norm** ``last_hidden_state``. The CausalLM simply passes this through, so ``hidden_states[-1]`` becomes + post-norm – a ~10x scale difference (std ~2 vs ~21) that breaks inference. + + This helper bypasses both mechanisms by hooking the last decoder layer directly, returning the raw pre-norm + output regardless of the transformers version. + """ + captured = {} + + def _hook(_module, _input, output): + 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: + forward_fn(**kwargs) + finally: + handle.remove() + return captured["hidden_states"] + + def encode_prompt_multiple_images( + self, + prompt: str | list[str], + device: torch.device | None = None, + images: list[Image.Image] | None = None, + max_sequence_length: int | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Encode prompts with inline tokens via the Qwen3-VL processor.""" + device = device or self._execution_device + template = self.prompt_template_encode["multiple_images"] + drop_idx = self.prompt_template_encode_start_idx["multiple_images"] + + prompt = [prompt] if isinstance(prompt, str) else prompt + prompt = [p.replace("\n", "<|vision_start|><|image_pad|><|vision_end|>") for p in prompt] + prompt = [template.format(p) for p in prompt] + + inputs = self.processor( + text=prompt, + images=images, + padding=True, + return_tensors="pt", + ).to(device) + + last_hidden_states = self._get_last_decoder_hidden_states(self.text_encoder, **inputs) + + prompt_embeds = last_hidden_states[:, drop_idx:] + prompt_embeds_mask = inputs["attention_mask"][:, drop_idx:] + + if max_sequence_length is not None and prompt_embeds.shape[1] > max_sequence_length: + prompt_embeds = prompt_embeds[:, -max_sequence_length:, :] + prompt_embeds_mask = prompt_embeds_mask[:, -max_sequence_length:] + + return prompt_embeds, prompt_embeds_mask + + def _pad_sequence(self, x: torch.Tensor, target_length: int) -> torch.Tensor: + current_length = x.shape[1] + if current_length >= target_length: + return x[:, -target_length:] + padding_length = target_length - current_length + if x.ndim >= 3: + padding = torch.zeros((x.shape[0], padding_length, *x.shape[2:]), dtype=x.dtype, device=x.device) + else: + padding = torch.zeros((x.shape[0], padding_length), dtype=x.dtype, device=x.device) + return torch.cat([x, padding], dim=1) + + def prepare_latents( + self, + batch_size: int, + num_channels_latents: int, + height: int, + width: int, + dtype: torch.dtype, + device: torch.device, + generator: torch.Generator | list[torch.Generator] | None, + reference_images: list[list[Image.Image]] | None = None, + latents: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, list[list[tuple[int, int, int]]]]: + """Prepare 6D padded latent tensor with target noise + reference image latents. + + Args: + latents: Optional pre-computed noise for the target slot. Shape ``(B, C, 1, H', W')`` where + ``H'`` and ``W'`` are the latent-space dimensions. When ``None``, random noise is sampled. + + Returns: + padded_latents: [B, max_patches, C, pt, ph, pw] target_mask: [B, max_patches] (True for target patches) + shape_list: per-sample list of (t, h, w) tuples for each component + """ + pt, ph, pw = self.transformer.config.patch_size + + all_patches = [] + all_target_masks = [] + all_shape_lists = [] + max_patches = 0 + + for i in range(batch_size): + sample_gen = generator[i] if isinstance(generator, list) else generator + + # Target noise + t_target = 1 + h_target = int(height) // self.vae_scale_factor_spatial + w_target = int(width) // self.vae_scale_factor_spatial + if latents is None: + noise_shape = (num_channels_latents, t_target, h_target, w_target) + noise_block = randn_tensor(noise_shape, generator=sample_gen, device=device, dtype=dtype) + else: + noise_block = latents[i].to(device=device, dtype=dtype) + + sample_items = [noise_block] + + # Reference images + if reference_images is not None and reference_images[i]: + for ref_img_pil in reference_images[i]: + ref_tensor = self.image_processor.preprocess(ref_img_pil).to(device=device, dtype=dtype) + ref_tensor = ref_tensor.unsqueeze(2) # [B, C, H, W] -> [B, C, 1, H, W] + + ref_latent = self.vae.encode(ref_tensor.to(self.vae.dtype)).latent_dist.mode() + ref_latent = ref_latent.to(dtype) + latents_mean = ( + torch.tensor(self.vae.config.latents_mean) + .view(1, -1, 1, 1, 1) + .to(ref_latent.device, ref_latent.dtype) + ) + latents_std = ( + torch.tensor(self.vae.config.latents_std) + .view(1, -1, 1, 1, 1) + .to(ref_latent.device, ref_latent.dtype) + ) + ref_latent = (ref_latent - latents_mean) / latents_std + ref_latent = ref_latent.squeeze(0) # [C, 1, H', W'] + sample_items.append(ref_latent) + + # Patchify each item and build shape_list + sample_patches = [] + sample_masks = [] + sample_shapes = [] + + for j, item in enumerate(sample_items): + c, t, h, w = item.shape + l_t, l_h, l_w = t // pt, h // ph, w // pw + sample_shapes.append((l_t, l_h, l_w)) + + patches = item.reshape(c, l_t, pt, l_h, ph, l_w, pw) + patches = patches.permute(1, 3, 5, 0, 2, 4, 6).reshape(-1, c, pt, ph, pw) + sample_patches.append(patches) + sample_masks.append(torch.full((patches.shape[0],), j == 0, device=device, dtype=torch.bool)) + + combined_patches = torch.cat(sample_patches, dim=0) + combined_masks = torch.cat(sample_masks, dim=0) + + all_patches.append(combined_patches) + all_target_masks.append(combined_masks) + all_shape_lists.append(sample_shapes) + max_patches = max(max_patches, combined_patches.shape[0]) + + # Pad to uniform size + padded_latents = torch.zeros( + (batch_size, max_patches, num_channels_latents, pt, ph, pw), device=device, dtype=dtype + ) + target_mask = torch.zeros((batch_size, max_patches), device=device, dtype=torch.bool) + + for i in range(batch_size): + n = all_patches[i].shape[0] + padded_latents[i, :n] = all_patches[i] + target_mask[i, :n] = all_target_masks[i] + + return padded_latents, target_mask, all_shape_lists + + # ------------------------------------------------------------------ + # Properties + # ------------------------------------------------------------------ + + @property + def guidance_scale(self) -> float: + return self._guidance_scale + + @property + def do_classifier_free_guidance(self) -> bool: + return self._guidance_scale > 1 + + @property + def num_timesteps(self) -> int: + return self._num_timesteps + + @property + def interrupt(self) -> bool: + return self._interrupt + + # ------------------------------------------------------------------ + # Forward pass + # ------------------------------------------------------------------ + + def check_inputs( + self, + prompt, + height, + width, + negative_prompt=None, + prompt_embeds=None, + negative_prompt_embeds=None, + callback_on_step_end_tensor_inputs=None, + ): + if height is not None and height % self.vae_scale_factor_spatial != 0: + raise ValueError(f"`height` must be divisible by {self.vae_scale_factor_spatial} but is {height}.") + if width is not None and width % self.vae_scale_factor_spatial != 0: + raise ValueError(f"`width` must be divisible by {self.vae_scale_factor_spatial} but is {width}.") + + if callback_on_step_end_tensor_inputs is not None and not all( + k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs + ): + raise ValueError( + f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found" + f" {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}" + ) + + if prompt is not None and prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to" + " only forward one of the two." + ) + elif prompt is None and prompt_embeds is None: + raise ValueError( + "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined." + ) + elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)): + raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") + + if negative_prompt is not None and negative_prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:" + f" {negative_prompt_embeds}. Please make sure to only forward one of the two." + ) + + @torch.no_grad() + @replace_example_docstring(EXAMPLE_DOC_STRING) + def __call__( + self, + images: list[Image.Image] | list[list[Image.Image]] | None = None, + prompt: str | list[str] = None, + height: int | None = None, + width: int | None = None, + num_inference_steps: int = 30, + timesteps: list[int] = None, + sigmas: list[float] = None, + guidance_scale: float = 4.0, + negative_prompt: str | list[str] | None = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + 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, + output_type: str | None = "pil", + return_dict: bool = True, + callback_on_step_end: Callable[[int, int, dict], None] + | PipelineCallback + | MultiPipelineCallbacks + | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], + max_sequence_length: int = 4096, + ): + r""" + Function invoked when calling the pipeline for generation. + + Args: + images (`list[Image.Image]` or `list[list[Image.Image]]`, *optional*): + Reference images for editing. Each image can have a different resolution. If a flat list is provided, + it is treated as one sample with multiple references. + prompt (`str` or `list[str]`, *optional*): + The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds` + instead. + height (`int`, *optional*): + The height in pixels of the generated image. If `None`, determined from the last reference image. + width (`int`, *optional*): + The width in pixels of the generated image. If `None`, determined from the last reference image. + num_inference_steps (`int`, *optional*, defaults to `30`): + The number of denoising steps. More denoising steps usually lead to a higher quality image at the + expense of slower inference. + timesteps (`list[int]`, *optional*): + Custom timesteps to use for the denoising process. If not defined, equal spacing is used. + sigmas (`list[float]`, *optional*): + Custom sigmas to use for the denoising process. + guidance_scale (`float`, *optional*, defaults to `4.0`): + Classifier-free guidance scale. Higher values encourage the model to generate images more aligned with + the `prompt` at the expense of lower image quality. + negative_prompt (`str` or `list[str]`, *optional*): + The prompt or prompts not to guide the image generation. If not defined, a blank prompt is used for + classifier-free guidance. + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): + One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) + to make generation deterministic. + latents (`torch.Tensor`, *optional*): + Pre-generated noisy latents to be used as inputs for image generation. + prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs. + prompt_embeds_mask (`torch.Tensor`, *optional*): + Attention mask for pre-generated text embeddings. + negative_prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated negative text embeddings. + negative_prompt_embeds_mask (`torch.Tensor`, *optional*): + Attention mask for pre-generated negative text embeddings. + output_type (`str`, *optional*, defaults to `"pil"`): + The output format of the generated image. Choose between `"pil"` (`PIL.Image.Image`), `"np"` + (`np.ndarray`), `"pt"` (`torch.Tensor`), or `"latent"` for raw latent output. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`JoyImageEditPlusPipelineOutput`] instead of a plain tuple. + callback_on_step_end (`Callable`, *optional*): + A function called at the end of each denoising step with arguments: the pipeline, step index, timestep, + and a dict of callback tensor inputs. + callback_on_step_end_tensor_inputs (`list[str]`, *optional*, defaults to `["latents"]`): + The list of tensor inputs for the `callback_on_step_end` function. + max_sequence_length (`int`, *optional*, defaults to `4096`): + Maximum sequence length for the text encoder. + + Examples: + + Returns: + [`JoyImageEditPlusPipelineOutput`] or `tuple`: + If `return_dict` is `True`, [`JoyImageEditPlusPipelineOutput`] is returned, otherwise a `tuple` is + returned where the first element is a list of generated images. + """ + # Normalize images input to List[List[Image]] + if images is not None: + if isinstance(images[0], Image.Image): + images = [images] # single sample + + self.check_inputs( + prompt=prompt, + height=height, + width=width, + negative_prompt=negative_prompt, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, + ) + + self._guidance_scale = guidance_scale + self._interrupt = False + + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + # Determine output resolution from last reference image if not specified + if height is None or width is None: + if images is not None and len(images[0]) > 0: + last_img = images[0][-1] + height, width = self.image_processor.get_default_height_width(last_img) + else: + height = height or 1024 + width = width or 1024 + + device = self._execution_device + + # Pre-process images: bucket-resize each reference image (matching original pipeline) + if images is not None: + processed_images = [] + for sample_imgs in images: + processed_sample = [] + for img in sample_imgs: + ref_h, ref_w = self.image_processor.get_default_height_width(img) + resize_img = self.image_processor.resize_center_crop(img, (ref_h, ref_w)) + processed_sample.append(resize_img) + processed_images.append(processed_sample) + images = processed_images + + # Construct prompts with tokens + prompt = [prompt] if isinstance(prompt, str) else prompt + if images is not None: + formatted_prompts = [] + for i in range(batch_size): + num_refs = len(images[i]) if i < len(images) else 0 + image_tags = "".join(["\n" for _ in range(num_refs)]) + p = prompt[i] if i < len(prompt) else prompt[0] + formatted_prompts.append(f"<|im_start|>user\n{image_tags}{p}<|im_end|>\n") + else: + formatted_prompts = [f"<|im_start|>user\n{p}<|im_end|>\n" for p in prompt] + + # Flatten all images for the processor + flattened_images = None + if images is not None: + flattened_images = [img for sublist in images for img in sublist] + + # Encode prompt + if prompt_embeds is None: + prompt_embeds, prompt_embeds_mask = self.encode_prompt_multiple_images( + prompt=formatted_prompts, + images=flattened_images, + device=device, + max_sequence_length=max_sequence_length, + ) + + # Encode negative prompt for CFG + if self.do_classifier_free_guidance: + if negative_prompt is None and negative_prompt_embeds is None: + neg_prompts = [] + for i in range(batch_size): + num_refs = len(images[i]) if images is not None and i < len(images) else 0 + image_tags = "".join(["\n" for _ in range(num_refs)]) + neg_prompts.append(f"<|im_start|>user\n{image_tags} <|im_end|>\n") + negative_prompt = neg_prompts + elif negative_prompt is not None and negative_prompt_embeds is None: + neg_list = [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt + neg_prompts = [] + for i in range(batch_size): + num_refs = len(images[i]) if images is not None and i < len(images) else 0 + image_tags = "".join(["\n" for _ in range(num_refs)]) + n = neg_list[i] if i < len(neg_list) else neg_list[0] + neg_prompts.append(f"<|im_start|>user\n{image_tags}{n}<|im_end|>\n") + negative_prompt = neg_prompts + + if negative_prompt_embeds is None: + neg_prompt_list = [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt + negative_prompt_embeds, negative_prompt_embeds_mask = self.encode_prompt_multiple_images( + prompt=neg_prompt_list, + images=flattened_images, + device=device, + max_sequence_length=max_sequence_length, + ) + + # Pad and concatenate [negative, positive] + max_seq_len = max(prompt_embeds.shape[1], negative_prompt_embeds.shape[1]) + prompt_embeds = torch.cat( + [ + self._pad_sequence(negative_prompt_embeds, max_seq_len), + self._pad_sequence(prompt_embeds, max_seq_len), + ] + ) + if prompt_embeds_mask is not None and negative_prompt_embeds_mask is not None: + prompt_embeds_mask = torch.cat( + [ + self._pad_sequence(negative_prompt_embeds_mask, max_seq_len), + self._pad_sequence(prompt_embeds_mask, max_seq_len), + ] + ) + + # Prepare timesteps + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, num_inference_steps, device, timesteps, sigmas + ) + + # Prepare latents (patchified) + num_channels_latents = self.transformer.config.in_channels + latents, target_mask, shape_list = self.prepare_latents( + batch_size=batch_size, + num_channels_latents=num_channels_latents, + height=height, + width=width, + dtype=prompt_embeds.dtype, + device=device, + generator=generator, + reference_images=images, + latents=latents, + ) + + # Denoising loop + num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order + self._num_timesteps = len(timesteps) + clean_reference_backup = latents.clone() + + with self.progress_bar(total=num_inference_steps) as progress_bar: + for i, t in enumerate(timesteps): + if self.interrupt: + continue + + # Restore reference patches + latents[~target_mask] = clean_reference_backup[~target_mask] + + model_input = latents + + # CFG expansion + if self.do_classifier_free_guidance: + model_input_cfg = torch.cat([model_input] * 2) + t_expand = t.repeat(model_input_cfg.shape[0]) + cfg_shape_list = shape_list * 2 + else: + model_input_cfg = model_input + t_expand = t.repeat(batch_size) + cfg_shape_list = shape_list + + # Transformer forward + noise_pred = self.transformer( + hidden_states=model_input_cfg, + timestep=t_expand, + encoder_hidden_states=prompt_embeds, + encoder_hidden_states_mask=prompt_embeds_mask, + shape_list=cfg_shape_list, + return_dict=False, + )[0] + + # CFG combination with norm rescaling + if self.do_classifier_free_guidance: + noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) + comb_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond) + cond_norm = torch.norm(noise_pred_text, dim=2, keepdim=True) + noise_norm = torch.norm(comb_pred, dim=2, keepdim=True) + noise_pred = comb_pred * (cond_norm / noise_norm.clamp_min(1e-6)) + + # Scheduler step + latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0].to( + dtype=prompt_embeds.dtype + ) + + if callback_on_step_end is not None: + callback_kwargs = {} + for k in callback_on_step_end_tensor_inputs: + callback_kwargs[k] = locals()[k] + callback_outputs = callback_on_step_end(self, i, t, callback_kwargs) + latents = callback_outputs.pop("latents", latents) + prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds) + + if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): + if progress_bar is not None: + progress_bar.update() + + # Post-processing: decode target latents + if output_type != "latent": + latents[~target_mask] = clean_reference_backup[~target_mask] + pt, ph, pw = self.transformer.config.patch_size + + image_list = [] + for b_idx in range(batch_size): + l_t, l_h, l_w = shape_list[b_idx][0] + target_len = l_t * l_h * l_w + + target_patches = latents[b_idx, :target_len] + c_lat = target_patches.shape[1] + video_latent = target_patches.reshape(l_t, l_h, l_w, c_lat, pt, ph, pw) + video_latent = video_latent.permute(3, 0, 4, 1, 5, 2, 6).reshape( + 1, c_lat, l_t * pt, l_h * ph, l_w * pw + ) + + latents_mean = ( + torch.tensor(self.vae.config.latents_mean) + .view(1, -1, 1, 1, 1) + .to(video_latent.device, video_latent.dtype) + ) + latents_std = ( + torch.tensor(self.vae.config.latents_std) + .view(1, -1, 1, 1, 1) + .to(video_latent.device, video_latent.dtype) + ) + video_latent = video_latent * latents_std + latents_mean + + sample_image = self.vae.decode(video_latent.to(self.vae.dtype), return_dict=False)[0] + # [1, C, T=1, H, W] -> [C, H, W] + sample_image = sample_image.float().squeeze(0).squeeze(1) + image_list.append(sample_image) + + image = torch.stack(image_list) # [B, C, H, W] + image = self.image_processor.postprocess(image, output_type=output_type) + else: + image = latents + + self.maybe_free_model_hooks() + + if not return_dict: + return (image,) + + return JoyImageEditPlusPipelineOutput(images=image) diff --git a/src/diffusers/pipelines/joyimage/pipeline_output.py b/src/diffusers/pipelines/joyimage/pipeline_output.py index 175dce3540d7..4ffb3e53a103 100644 --- a/src/diffusers/pipelines/joyimage/pipeline_output.py +++ b/src/diffusers/pipelines/joyimage/pipeline_output.py @@ -14,3 +14,12 @@ class JoyImageEditPipelineOutput(BaseOutput): """ images: Union[List[PIL.Image.Image], np.ndarray] + + +@dataclass +class JoyImageEditPlusPipelineOutput(BaseOutput): + """ + Output class for JoyImage Edit Plus multi-image editing pipelines. + """ + + images: Union[List[PIL.Image.Image], np.ndarray] diff --git a/src/diffusers/pipelines/kolors/tokenizer.py b/src/diffusers/pipelines/kolors/tokenizer.py index c0a1831bdeee..5fc475ddc079 100644 --- a/src/diffusers/pipelines/kolors/tokenizer.py +++ b/src/diffusers/pipelines/kolors/tokenizer.py @@ -16,14 +16,22 @@ import os import re -from sentencepiece import SentencePieceProcessor from transformers import PreTrainedTokenizer from transformers.tokenization_utils_base import BatchEncoding, EncodedInput from transformers.utils import PaddingStrategy +from ...utils import is_sentencepiece_available + class SPTokenizer: def __init__(self, model_path: str): + if not is_sentencepiece_available(): + raise ImportError( + "`SPTokenizer` requires the `sentencepiece` library but it was not found in your environment. You can " + "install it with `pip install sentencepiece`." + ) + from sentencepiece import SentencePieceProcessor + # reload tokenizer assert os.path.isfile(model_path), model_path self.sp_model = SentencePieceProcessor(model_file=model_path) diff --git a/src/diffusers/pipelines/ledits_pp/pipeline_leditspp_stable_diffusion.py b/src/diffusers/pipelines/ledits_pp/pipeline_leditspp_stable_diffusion.py index 864f9feeb5aa..9048b56b01f4 100644 --- a/src/diffusers/pipelines/ledits_pp/pipeline_leditspp_stable_diffusion.py +++ b/src/diffusers/pipelines/ledits_pp/pipeline_leditspp_stable_diffusion.py @@ -51,7 +51,7 @@ >>> pipe = LEditsPPPipelineStableDiffusion.from_pretrained( ... "stable-diffusion-v1-5/stable-diffusion-v1-5", variant="fp16", torch_dtype=torch.float16 ... ) - >>> pipe.enable_vae_tiling() + >>> pipe.vae.enable_tiling() >>> pipe = pipe.to("cuda") >>> img_url = "https://www.aiml.informatik.tu-darmstadt.de/people/mbrack/cherry_blossom.png" @@ -717,59 +717,6 @@ def clip_skip(self): def cross_attention_kwargs(self): return self._cross_attention_kwargs - def enable_vae_slicing(self): - r""" - Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to - compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. - """ - depr_message = f"Calling `enable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_slicing()`." - deprecate( - "enable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.enable_slicing() - - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_slicing()`." - deprecate( - "disable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.disable_slicing() - - def enable_vae_tiling(self): - r""" - Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to - compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow - processing larger images. - """ - depr_message = f"Calling `enable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_tiling()`." - deprecate( - "enable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.enable_tiling() - - def disable_vae_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_tiling()`." - deprecate( - "disable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.disable_tiling() - @torch.no_grad() @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( diff --git a/src/diffusers/pipelines/ledits_pp/pipeline_leditspp_stable_diffusion_xl.py b/src/diffusers/pipelines/ledits_pp/pipeline_leditspp_stable_diffusion_xl.py index 70a61fab1be2..c2e6f3caaf47 100644 --- a/src/diffusers/pipelines/ledits_pp/pipeline_leditspp_stable_diffusion_xl.py +++ b/src/diffusers/pipelines/ledits_pp/pipeline_leditspp_stable_diffusion_xl.py @@ -78,7 +78,7 @@ >>> pipe = LEditsPPPipelineStableDiffusionXL.from_pretrained( ... "stabilityai/stable-diffusion-xl-base-1.0", variant="fp16", torch_dtype=torch.float16 ... ) - >>> pipe.enable_vae_tiling() + >>> pipe.vae.enable_tiling() >>> pipe = pipe.to("cuda") >>> img_url = "https://www.aiml.informatik.tu-darmstadt.de/people/mbrack/tennis.jpg" @@ -755,59 +755,6 @@ def denoising_end(self): def num_timesteps(self): return self._num_timesteps - def enable_vae_slicing(self): - r""" - Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to - compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. - """ - depr_message = f"Calling `enable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_slicing()`." - deprecate( - "enable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.enable_slicing() - - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_slicing()`." - deprecate( - "disable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.disable_slicing() - - def enable_vae_tiling(self): - r""" - Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to - compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow - processing larger images. - """ - depr_message = f"Calling `enable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_tiling()`." - deprecate( - "enable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.enable_tiling() - - def disable_vae_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_tiling()`." - deprecate( - "disable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.disable_tiling() - # Copied from diffusers.pipelines.ledits_pp.pipeline_leditspp_stable_diffusion.LEditsPPPipelineStableDiffusion.prepare_unet def prepare_unet(self, attention_store, PnP: bool = False): attn_procs = {} diff --git a/src/diffusers/pipelines/ltx/pipeline_ltx_latent_upsample.py b/src/diffusers/pipelines/ltx/pipeline_ltx_latent_upsample.py index 315dcc04cb30..fbf98ce976d9 100644 --- a/src/diffusers/pipelines/ltx/pipeline_ltx_latent_upsample.py +++ b/src/diffusers/pipelines/ltx/pipeline_ltx_latent_upsample.py @@ -16,7 +16,7 @@ from ...image_processor import PipelineImageInput from ...models import AutoencoderKLLTXVideo -from ...utils import deprecate, get_logger +from ...utils import get_logger from ...utils.torch_utils import randn_tensor from ...video_processor import VideoProcessor from ..pipeline_utils import DiffusionPipeline @@ -173,59 +173,6 @@ def _denormalize_latents( latents = latents * latents_std / scaling_factor + latents_mean return latents - def enable_vae_slicing(self): - r""" - Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to - compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. - """ - depr_message = f"Calling `enable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_slicing()`." - deprecate( - "enable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.enable_slicing() - - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_slicing()`." - deprecate( - "disable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.disable_slicing() - - def enable_vae_tiling(self): - r""" - Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to - compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow - processing larger images. - """ - depr_message = f"Calling `enable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_tiling()`." - deprecate( - "enable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.enable_tiling() - - def disable_vae_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_tiling()`." - deprecate( - "disable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.disable_tiling() - def check_inputs(self, video, height, width, latents, tone_map_compression_ratio): if height % self.vae_spatial_compression_ratio != 0 or width % self.vae_spatial_compression_ratio != 0: raise ValueError(f"`height` and `width` have to be divisible by 32 but are {height} and {width}.") diff --git a/src/diffusers/pipelines/lumina2/pipeline_lumina2.py b/src/diffusers/pipelines/lumina2/pipeline_lumina2.py index 8a7a8925a925..2129bfeb8fbf 100644 --- a/src/diffusers/pipelines/lumina2/pipeline_lumina2.py +++ b/src/diffusers/pipelines/lumina2/pipeline_lumina2.py @@ -428,59 +428,6 @@ def check_inputs( if max_sequence_length is not None and max_sequence_length > 512: raise ValueError(f"`max_sequence_length` cannot be greater than 512 but is {max_sequence_length}") - def enable_vae_slicing(self): - r""" - Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to - compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. - """ - depr_message = f"Calling `enable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_slicing()`." - deprecate( - "enable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.enable_slicing() - - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_slicing()`." - deprecate( - "disable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.disable_slicing() - - def enable_vae_tiling(self): - r""" - Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to - compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow - processing larger images. - """ - depr_message = f"Calling `enable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_tiling()`." - deprecate( - "enable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.enable_tiling() - - def disable_vae_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_tiling()`." - deprecate( - "disable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.disable_tiling() - def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None): # VAE applies 8x compression on images but we must also account for packing which requires # latent height and width to be divisible by 2. diff --git a/src/diffusers/pipelines/mochi/pipeline_mochi.py b/src/diffusers/pipelines/mochi/pipeline_mochi.py index 0e791b5f6b20..7108d4ff774f 100644 --- a/src/diffusers/pipelines/mochi/pipeline_mochi.py +++ b/src/diffusers/pipelines/mochi/pipeline_mochi.py @@ -23,7 +23,7 @@ from ...loaders import Mochi1LoraLoaderMixin from ...models import AutoencoderKLMochi, MochiTransformer3DModel from ...schedulers import FlowMatchEulerDiscreteScheduler -from ...utils import deprecate, is_torch_xla_available, logging, replace_example_docstring +from ...utils import is_torch_xla_available, logging, replace_example_docstring from ...utils.torch_utils import randn_tensor from ...video_processor import VideoProcessor from ..pipeline_utils import DiffusionPipeline @@ -49,7 +49,7 @@ >>> pipe = MochiPipeline.from_pretrained("genmo/mochi-1-preview", torch_dtype=torch.bfloat16) >>> pipe.enable_model_cpu_offload() - >>> pipe.enable_vae_tiling() + >>> pipe.vae.enable_tiling() >>> prompt = "Close-up of a chameleon's eye, with its scaly skin changing color. Ultra high resolution 4k." >>> frames = pipe(prompt, num_inference_steps=28, guidance_scale=3.5).frames[0] >>> export_to_video(frames, "mochi.mp4") @@ -387,59 +387,6 @@ def check_inputs( f" {negative_prompt_attention_mask.shape}." ) - def enable_vae_slicing(self): - r""" - Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to - compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. - """ - depr_message = f"Calling `enable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_slicing()`." - deprecate( - "enable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.enable_slicing() - - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_slicing()`." - deprecate( - "disable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.disable_slicing() - - def enable_vae_tiling(self): - r""" - Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to - compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow - processing larger images. - """ - depr_message = f"Calling `enable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_tiling()`." - deprecate( - "enable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.enable_tiling() - - def disable_vae_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_tiling()`." - deprecate( - "disable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.disable_tiling() - def prepare_latents( self, batch_size, diff --git a/src/diffusers/pipelines/omnigen/pipeline_omnigen.py b/src/diffusers/pipelines/omnigen/pipeline_omnigen.py index bdf4e30c6619..6e5db93d1f35 100644 --- a/src/diffusers/pipelines/omnigen/pipeline_omnigen.py +++ b/src/diffusers/pipelines/omnigen/pipeline_omnigen.py @@ -23,7 +23,7 @@ from ...models.autoencoders import AutoencoderKL from ...models.transformers import OmniGenTransformer2DModel from ...schedulers import FlowMatchEulerDiscreteScheduler -from ...utils import deprecate, is_torch_xla_available, is_torchvision_available, logging, replace_example_docstring +from ...utils import is_torch_xla_available, is_torchvision_available, logging, replace_example_docstring from ...utils.torch_utils import randn_tensor from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput @@ -230,59 +230,6 @@ def check_inputs( f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}" ) - def enable_vae_slicing(self): - r""" - Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to - compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. - """ - depr_message = f"Calling `enable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_slicing()`." - deprecate( - "enable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.enable_slicing() - - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_slicing()`." - deprecate( - "disable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.disable_slicing() - - def enable_vae_tiling(self): - r""" - Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to - compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow - processing larger images. - """ - depr_message = f"Calling `enable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_tiling()`." - deprecate( - "enable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.enable_tiling() - - def disable_vae_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_tiling()`." - deprecate( - "disable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.disable_tiling() - # Copied from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3.StableDiffusion3Pipeline.prepare_latents def prepare_latents( self, diff --git a/src/diffusers/pipelines/pag/pipeline_pag_sana.py b/src/diffusers/pipelines/pag/pipeline_pag_sana.py index 71861f366477..f621255325a0 100644 --- a/src/diffusers/pipelines/pag/pipeline_pag_sana.py +++ b/src/diffusers/pipelines/pag/pipeline_pag_sana.py @@ -29,7 +29,6 @@ from ...schedulers import FlowMatchEulerDiscreteScheduler from ...utils import ( BACKENDS_MAPPING, - deprecate, is_bs4_available, is_ftfy_available, is_torch_xla_available, @@ -186,59 +185,6 @@ def __init__( pag_attn_processors=(PAGCFGSanaLinearAttnProcessor2_0(), PAGIdentitySanaLinearAttnProcessor2_0()), ) - def enable_vae_slicing(self): - r""" - Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to - compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. - """ - depr_message = f"Calling `enable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_slicing()`." - deprecate( - "enable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.enable_slicing() - - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_slicing()`." - deprecate( - "disable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.disable_slicing() - - def enable_vae_tiling(self): - r""" - Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to - compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow - processing larger images. - """ - depr_message = f"Calling `enable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_tiling()`." - deprecate( - "enable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.enable_tiling() - - def disable_vae_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_tiling()`." - deprecate( - "disable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.disable_tiling() - def encode_prompt( self, prompt: str | list[str], diff --git a/src/diffusers/pipelines/pipeline_loading_utils.py b/src/diffusers/pipelines/pipeline_loading_utils.py index d695f5e7284d..90cbffc5b69d 100644 --- a/src/diffusers/pipelines/pipeline_loading_utils.py +++ b/src/diffusers/pipelines/pipeline_loading_utils.py @@ -69,6 +69,19 @@ TRANSFORMERS_DUMMY_MODULES_FOLDER = "transformers.utils" CONNECTED_PIPES_KEYS = ["prior"] +# Auxiliary (non-weight) files a transformers component saves next to its weights. Repos with a flat, +# transformers-style layout host a component's files at the repo root instead of in a subfolder, where the +# folder-based allow patterns of `DiffusionPipeline.download` would miss them. Root-hosted weights and +# `config.json` are matched by their own patterns, so only these auxiliary filenames need listing. +# Currently the set needed by DiffusionGemma — extend as new flat-layout pipelines require it. +TRANSFORMERS_COMPONENT_AUX_FILES = [ + "chat_template.jinja", + "generation_config.json", + "processor_config.json", + "tokenizer.json", + "tokenizer_config.json", +] + logger = logging.get_logger(__name__) LOADABLE_CLASSES = { @@ -136,6 +149,8 @@ def is_safetensors_compatible(filenames, passed_components=None, folder_names=No ) passed_components = passed_components or [] + # only weight files matter for safetensors compatibility + filenames = filter_model_files(filenames) if folder_names: filenames = {f for f in filenames if os.path.split(f)[0] in folder_names} diff --git a/src/diffusers/pipelines/pipeline_utils.py b/src/diffusers/pipelines/pipeline_utils.py index 1fa4db90d995..7b653caa0657 100644 --- a/src/diffusers/pipelines/pipeline_utils.py +++ b/src/diffusers/pipelines/pipeline_utils.py @@ -58,7 +58,6 @@ PushToHubMixin, _get_detailed_type, _is_valid_type, - deprecate, is_accelerate_available, is_accelerate_version, is_bitsandbytes_version, @@ -82,6 +81,7 @@ CONNECTED_PIPES_KEYS, CUSTOM_PIPELINE_FILE_NAME, LOADABLE_CLASSES, + TRANSFORMERS_COMPONENT_AUX_FILES, _download_dduf_file, _fetch_class_library_tuple, _get_custom_components_and_folders, @@ -1751,6 +1751,11 @@ def download(cls, pretrained_model_name, **kwargs) -> str | os.PathLike: p for p in allow_patterns if not (len(p.split("/")) == 2 and p.split("/")[0] in passed_components) ] + # Repos with a flat, transformers-style layout host a component's files at the repo root instead of + # in a subfolder, where the folder-based allow patterns above miss its auxiliary files (root-hosted + # weights are already included via `model_filenames`, root `config.json` via `CONFIG_NAME`). + allow_patterns += TRANSFORMERS_COMPONENT_AUX_FILES + if pipeline_class._load_connected_pipes: allow_patterns.append("README.md") @@ -2258,59 +2263,6 @@ class StableDiffusionMixin: Helper for DiffusionPipeline with vae and unet.(mainly for LDM such as stable diffusion) """ - def enable_vae_slicing(self): - r""" - Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to - compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. - """ - depr_message = f"Calling `enable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_slicing()`." - deprecate( - "enable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.enable_slicing() - - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_slicing()`." - deprecate( - "disable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.disable_slicing() - - def enable_vae_tiling(self): - r""" - Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to - compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow - processing larger images. - """ - depr_message = f"Calling `enable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_tiling()`." - deprecate( - "enable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.enable_tiling() - - def disable_vae_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_tiling()`." - deprecate( - "disable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.disable_tiling() - def enable_freeu(self, s1: float, s2: float, b1: float, b2: float): r"""Enables the FreeU mechanism as in https://huggingface.co/papers/2309.11497. diff --git a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage.py b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage.py index 7806db6546d4..1da0518a4f65 100644 --- a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage.py +++ b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage.py @@ -23,7 +23,7 @@ from ...loaders import QwenImageLoraLoaderMixin from ...models import AutoencoderKLQwenImage, QwenImageTransformer2DModel from ...schedulers import FlowMatchEulerDiscreteScheduler -from ...utils import deprecate, is_torch_xla_available, logging, replace_example_docstring +from ...utils import is_torch_xla_available, logging, replace_example_docstring from ...utils.torch_utils import randn_tensor from ..pipeline_utils import DiffusionPipeline from .pipeline_output import QwenImagePipelineOutput @@ -353,59 +353,6 @@ def _unpack_latents(latents, height, width, vae_scale_factor): return latents - def enable_vae_slicing(self): - r""" - Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to - compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. - """ - depr_message = f"Calling `enable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_slicing()`." - deprecate( - "enable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.enable_slicing() - - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_slicing()`." - deprecate( - "disable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.disable_slicing() - - def enable_vae_tiling(self): - r""" - Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to - compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow - processing larger images. - """ - depr_message = f"Calling `enable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_tiling()`." - deprecate( - "enable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.enable_tiling() - - def disable_vae_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_tiling()`." - deprecate( - "disable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.disable_tiling() - def prepare_latents( self, batch_size, diff --git a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_controlnet.py b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_controlnet.py index 1b0a460f95c7..f946fdf27d00 100644 --- a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_controlnet.py +++ b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_controlnet.py @@ -24,7 +24,7 @@ from ...models import AutoencoderKLQwenImage, QwenImageTransformer2DModel from ...models.controlnets.controlnet_qwenimage import QwenImageControlNetModel, QwenImageMultiControlNetModel from ...schedulers import FlowMatchEulerDiscreteScheduler -from ...utils import deprecate, is_torch_xla_available, logging, replace_example_docstring +from ...utils import is_torch_xla_available, logging, replace_example_docstring from ...utils.torch_utils import randn_tensor from ..pipeline_utils import DiffusionPipeline from .pipeline_output import QwenImagePipelineOutput @@ -420,59 +420,6 @@ def _unpack_latents(latents, height, width, vae_scale_factor): return latents - def enable_vae_slicing(self): - r""" - Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to - compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. - """ - depr_message = f"Calling `enable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_slicing()`." - deprecate( - "enable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.enable_slicing() - - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_slicing()`." - deprecate( - "disable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.disable_slicing() - - def enable_vae_tiling(self): - r""" - Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to - compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow - processing larger images. - """ - depr_message = f"Calling `enable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_tiling()`." - deprecate( - "enable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.enable_tiling() - - def disable_vae_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_tiling()`." - deprecate( - "disable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.disable_tiling() - # Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage.QwenImagePipeline.prepare_latents def prepare_latents( self, diff --git a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_controlnet_inpaint.py b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_controlnet_inpaint.py index 75a11c183542..97f510a6dbf4 100644 --- a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_controlnet_inpaint.py +++ b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_controlnet_inpaint.py @@ -402,35 +402,6 @@ def _unpack_latents(latents, height, width, vae_scale_factor): return latents - def enable_vae_slicing(self): - r""" - Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to - compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. - """ - self.vae.enable_slicing() - - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - self.vae.disable_slicing() - - def enable_vae_tiling(self): - r""" - Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to - compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow - processing larger images. - """ - self.vae.enable_tiling() - - def disable_vae_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to - computing decoding in one step. - """ - self.vae.disable_tiling() - # Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage.QwenImagePipeline.prepare_latents def prepare_latents( self, diff --git a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit.py b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit.py index d356d932613d..85abb815cf23 100644 --- a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit.py +++ b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit.py @@ -24,7 +24,7 @@ from ...loaders import QwenImageLoraLoaderMixin from ...models import AutoencoderKLQwenImage, QwenImageTransformer2DModel from ...schedulers import FlowMatchEulerDiscreteScheduler -from ...utils import deprecate, is_torch_xla_available, logging, replace_example_docstring +from ...utils import is_torch_xla_available, logging, replace_example_docstring from ...utils.torch_utils import randn_tensor from ..pipeline_utils import DiffusionPipeline from .pipeline_output import QwenImagePipelineOutput @@ -426,59 +426,6 @@ def _encode_vae_image(self, image: torch.Tensor, generator: torch.Generator): return image_latents - def enable_vae_slicing(self): - r""" - Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to - compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. - """ - depr_message = f"Calling `enable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_slicing()`." - deprecate( - "enable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.enable_slicing() - - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_slicing()`." - deprecate( - "disable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.disable_slicing() - - def enable_vae_tiling(self): - r""" - Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to - compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow - processing larger images. - """ - depr_message = f"Calling `enable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_tiling()`." - deprecate( - "enable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.enable_tiling() - - def disable_vae_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_tiling()`." - deprecate( - "disable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.disable_tiling() - def prepare_latents( self, image, diff --git a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit_inpaint.py b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit_inpaint.py index fe4fdb7dfda3..57d1fdaaf99f 100644 --- a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit_inpaint.py +++ b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_edit_inpaint.py @@ -25,7 +25,7 @@ from ...loaders import QwenImageLoraLoaderMixin from ...models import AutoencoderKLQwenImage, QwenImageTransformer2DModel from ...schedulers import FlowMatchEulerDiscreteScheduler -from ...utils import deprecate, is_torch_xla_available, logging, replace_example_docstring +from ...utils import is_torch_xla_available, logging, replace_example_docstring from ...utils.torch_utils import randn_tensor from ..pipeline_utils import DiffusionPipeline from .pipeline_output import QwenImagePipelineOutput @@ -472,59 +472,6 @@ def get_timesteps(self, num_inference_steps, strength, device): return timesteps, num_inference_steps - t_start - def enable_vae_slicing(self): - r""" - Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to - compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. - """ - depr_message = f"Calling `enable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_slicing()`." - deprecate( - "enable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.enable_slicing() - - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_slicing()`." - deprecate( - "disable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.disable_slicing() - - def enable_vae_tiling(self): - r""" - Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to - compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow - processing larger images. - """ - depr_message = f"Calling `enable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_tiling()`." - deprecate( - "enable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.enable_tiling() - - def disable_vae_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_tiling()`." - deprecate( - "disable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.disable_tiling() - # Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage_inpaint.QwenImageInpaintPipeline.prepare_latents def prepare_latents( self, diff --git a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_img2img.py b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_img2img.py index 95d7d80c861c..9b9af83737e5 100644 --- a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_img2img.py +++ b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_img2img.py @@ -9,7 +9,7 @@ from ...loaders import QwenImageLoraLoaderMixin from ...models import AutoencoderKLQwenImage, QwenImageTransformer2DModel from ...schedulers import FlowMatchEulerDiscreteScheduler -from ...utils import deprecate, is_torch_xla_available, logging, replace_example_docstring +from ...utils import is_torch_xla_available, logging, replace_example_docstring from ...utils.torch_utils import randn_tensor from ..pipeline_utils import DiffusionPipeline from .pipeline_output import QwenImagePipelineOutput @@ -402,59 +402,6 @@ def _unpack_latents(latents, height, width, vae_scale_factor): return latents - def enable_vae_slicing(self): - r""" - Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to - compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. - """ - depr_message = f"Calling `enable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_slicing()`." - deprecate( - "enable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.enable_slicing() - - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_slicing()`." - deprecate( - "disable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.disable_slicing() - - def enable_vae_tiling(self): - r""" - Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to - compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow - processing larger images. - """ - depr_message = f"Calling `enable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_tiling()`." - deprecate( - "enable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.enable_tiling() - - def disable_vae_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_tiling()`." - deprecate( - "disable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.disable_tiling() - def prepare_latents( self, image, diff --git a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_inpaint.py b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_inpaint.py index 01882bddefa7..3d5f0040932a 100644 --- a/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_inpaint.py +++ b/src/diffusers/pipelines/qwenimage/pipeline_qwenimage_inpaint.py @@ -10,7 +10,7 @@ from ...loaders import QwenImageLoraLoaderMixin from ...models import AutoencoderKLQwenImage, QwenImageTransformer2DModel from ...schedulers import FlowMatchEulerDiscreteScheduler -from ...utils import deprecate, is_torch_xla_available, logging, replace_example_docstring +from ...utils import is_torch_xla_available, logging, replace_example_docstring from ...utils.torch_utils import randn_tensor from ..pipeline_utils import DiffusionPipeline from .pipeline_output import QwenImagePipelineOutput @@ -429,59 +429,6 @@ def _unpack_latents(latents, height, width, vae_scale_factor): return latents - def enable_vae_slicing(self): - r""" - Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to - compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. - """ - depr_message = f"Calling `enable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_slicing()`." - deprecate( - "enable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.enable_slicing() - - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_slicing()`." - deprecate( - "disable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.disable_slicing() - - def enable_vae_tiling(self): - r""" - Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to - compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow - processing larger images. - """ - depr_message = f"Calling `enable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_tiling()`." - deprecate( - "enable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.enable_tiling() - - def disable_vae_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_tiling()`." - deprecate( - "disable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.disable_tiling() - def prepare_latents( self, image, diff --git a/src/diffusers/pipelines/sana/pipeline_sana.py b/src/diffusers/pipelines/sana/pipeline_sana.py index 17e0be9ba7aa..553e45a628d9 100644 --- a/src/diffusers/pipelines/sana/pipeline_sana.py +++ b/src/diffusers/pipelines/sana/pipeline_sana.py @@ -30,7 +30,6 @@ from ...utils import ( BACKENDS_MAPPING, USE_PEFT_BACKEND, - deprecate, is_bs4_available, is_ftfy_available, is_torch_xla_available, @@ -220,59 +219,6 @@ def __init__( ) self.image_processor = PixArtImageProcessor(vae_scale_factor=self.vae_scale_factor) - def enable_vae_slicing(self): - r""" - Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to - compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. - """ - depr_message = f"Calling `enable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_slicing()`." - deprecate( - "enable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.enable_slicing() - - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_slicing()`." - deprecate( - "disable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.disable_slicing() - - def enable_vae_tiling(self): - r""" - Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to - compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow - processing larger images. - """ - depr_message = f"Calling `enable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_tiling()`." - deprecate( - "enable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.enable_tiling() - - def disable_vae_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_tiling()`." - deprecate( - "disable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.disable_tiling() - def _get_gemma_prompt_embeds( self, prompt: str | list[str], diff --git a/src/diffusers/pipelines/sana/pipeline_sana_controlnet.py b/src/diffusers/pipelines/sana/pipeline_sana_controlnet.py index d976c7035d9d..de1910f68192 100644 --- a/src/diffusers/pipelines/sana/pipeline_sana_controlnet.py +++ b/src/diffusers/pipelines/sana/pipeline_sana_controlnet.py @@ -30,7 +30,6 @@ from ...utils import ( BACKENDS_MAPPING, USE_PEFT_BACKEND, - deprecate, is_bs4_available, is_ftfy_available, is_torch_xla_available, @@ -233,59 +232,6 @@ def __init__( ) self.image_processor = PixArtImageProcessor(vae_scale_factor=self.vae_scale_factor) - def enable_vae_slicing(self): - r""" - Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to - compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. - """ - depr_message = f"Calling `enable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_slicing()`." - deprecate( - "enable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.enable_slicing() - - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_slicing()`." - deprecate( - "disable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.disable_slicing() - - def enable_vae_tiling(self): - r""" - Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to - compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow - processing larger images. - """ - depr_message = f"Calling `enable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_tiling()`." - deprecate( - "enable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.enable_tiling() - - def disable_vae_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_tiling()`." - deprecate( - "disable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.disable_tiling() - # Copied from diffusers.pipelines.sana.pipeline_sana.SanaPipeline._get_gemma_prompt_embeds def _get_gemma_prompt_embeds( self, diff --git a/src/diffusers/pipelines/sana/pipeline_sana_sprint.py b/src/diffusers/pipelines/sana/pipeline_sana_sprint.py index c85f05275f51..812441d8e462 100644 --- a/src/diffusers/pipelines/sana/pipeline_sana_sprint.py +++ b/src/diffusers/pipelines/sana/pipeline_sana_sprint.py @@ -30,7 +30,6 @@ from ...utils import ( BACKENDS_MAPPING, USE_PEFT_BACKEND, - deprecate, is_bs4_available, is_ftfy_available, is_torch_xla_available, @@ -171,59 +170,6 @@ def __init__( ) self.image_processor = PixArtImageProcessor(vae_scale_factor=self.vae_scale_factor) - def enable_vae_slicing(self): - r""" - Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to - compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. - """ - depr_message = f"Calling `enable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_slicing()`." - deprecate( - "enable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.enable_slicing() - - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_slicing()`." - deprecate( - "disable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.disable_slicing() - - def enable_vae_tiling(self): - r""" - Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to - compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow - processing larger images. - """ - depr_message = f"Calling `enable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_tiling()`." - deprecate( - "enable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.enable_tiling() - - def disable_vae_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_tiling()`." - deprecate( - "disable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.disable_tiling() - # Copied from diffusers.pipelines.sana.pipeline_sana.SanaPipeline._get_gemma_prompt_embeds def _get_gemma_prompt_embeds( self, diff --git a/src/diffusers/pipelines/sana/pipeline_sana_sprint_img2img.py b/src/diffusers/pipelines/sana/pipeline_sana_sprint_img2img.py index b3bd7b776d81..e149e0c597b2 100644 --- a/src/diffusers/pipelines/sana/pipeline_sana_sprint_img2img.py +++ b/src/diffusers/pipelines/sana/pipeline_sana_sprint_img2img.py @@ -31,7 +31,6 @@ from ...utils import ( BACKENDS_MAPPING, USE_PEFT_BACKEND, - deprecate, is_bs4_available, is_ftfy_available, is_torch_xla_available, @@ -178,62 +177,6 @@ def __init__( ) self.image_processor = PixArtImageProcessor(vae_scale_factor=self.vae_scale_factor) - # Copied from diffusers.pipelines.sana.pipeline_sana.SanaPipeline.enable_vae_slicing - def enable_vae_slicing(self): - r""" - Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to - compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. - """ - depr_message = f"Calling `enable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_slicing()`." - deprecate( - "enable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.enable_slicing() - - # Copied from diffusers.pipelines.sana.pipeline_sana.SanaPipeline.disable_vae_slicing - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_slicing()`." - deprecate( - "disable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.disable_slicing() - - # Copied from diffusers.pipelines.sana.pipeline_sana.SanaPipeline.enable_vae_tiling - def enable_vae_tiling(self): - r""" - Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to - compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow - processing larger images. - """ - depr_message = f"Calling `enable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_tiling()`." - deprecate( - "enable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.enable_tiling() - - def disable_vae_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_tiling()`." - deprecate( - "disable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.disable_tiling() - # Copied from diffusers.pipelines.sana.pipeline_sana.SanaPipeline._get_gemma_prompt_embeds def _get_gemma_prompt_embeds( self, diff --git a/src/diffusers/pipelines/stable_audio/pipeline_stable_audio.py b/src/diffusers/pipelines/stable_audio/pipeline_stable_audio.py index 351c8b65de0e..475f4032edab 100644 --- a/src/diffusers/pipelines/stable_audio/pipeline_stable_audio.py +++ b/src/diffusers/pipelines/stable_audio/pipeline_stable_audio.py @@ -25,7 +25,7 @@ from ...models import AutoencoderOobleck, StableAudioDiTModel from ...models.embeddings import get_1d_rotary_pos_embed from ...schedulers import EDMDPMSolverMultistepScheduler -from ...utils import deprecate, is_torch_xla_available, logging, replace_example_docstring +from ...utils import is_torch_xla_available, logging, replace_example_docstring from ...utils.torch_utils import randn_tensor from ..pipeline_utils import AudioPipelineOutput, DiffusionPipeline from .modeling_stable_audio import StableAudioProjectionModel @@ -124,34 +124,6 @@ def __init__( ) self.rotary_embed_dim = self.transformer.config.attention_head_dim // 2 - # Copied from diffusers.pipelines.pipeline_utils.StableDiffusionMixin.enable_vae_slicing - def enable_vae_slicing(self): - r""" - Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to - compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. - """ - depr_message = f"Calling `enable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_slicing()`." - deprecate( - "enable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.enable_slicing() - - # Copied from diffusers.pipelines.pipeline_utils.StableDiffusionMixin.disable_vae_slicing - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_slicing()`." - deprecate( - "disable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.disable_slicing() - def encode_prompt( self, prompt, diff --git a/src/diffusers/pipelines/visualcloze/pipeline_visualcloze_generation.py b/src/diffusers/pipelines/visualcloze/pipeline_visualcloze_generation.py index dd5d0603d6d0..17aa366ebacd 100644 --- a/src/diffusers/pipelines/visualcloze/pipeline_visualcloze_generation.py +++ b/src/diffusers/pipelines/visualcloze/pipeline_visualcloze_generation.py @@ -24,7 +24,6 @@ from ...schedulers import FlowMatchEulerDiscreteScheduler from ...utils import ( USE_PEFT_BACKEND, - deprecate, is_torch_xla_available, logging, replace_example_docstring, @@ -520,59 +519,6 @@ def _unpack_latents(latents, sizes, vae_scale_factor): return unpacked_latents - def enable_vae_slicing(self): - r""" - Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to - compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. - """ - depr_message = f"Calling `enable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_slicing()`." - deprecate( - "enable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.enable_slicing() - - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_slicing()`." - deprecate( - "disable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.disable_slicing() - - def enable_vae_tiling(self): - r""" - Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to - compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow - processing larger images. - """ - depr_message = f"Calling `enable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_tiling()`." - deprecate( - "enable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.enable_tiling() - - def disable_vae_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = f"Calling `disable_vae_tiling()` on a `{self.__class__.__name__}` is deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_tiling()`." - deprecate( - "disable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.disable_tiling() - def _prepare_latents(self, image, mask, gen, vae_scale_factor, device, dtype): """Helper function to prepare latents for a single batch.""" # Concatenate images and masks along width dimension diff --git a/src/diffusers/schedulers/scheduling_flow_match_euler_discrete.py b/src/diffusers/schedulers/scheduling_flow_match_euler_discrete.py index 7b207f782079..0e4f5c6a1f97 100644 --- a/src/diffusers/schedulers/scheduling_flow_match_euler_discrete.py +++ b/src/diffusers/schedulers/scheduling_flow_match_euler_discrete.py @@ -364,10 +364,7 @@ def set_timesteps( # 5. Convert sigmas and timesteps to tensors and move to specified device sigmas = torch.from_numpy(sigmas).to(dtype=torch.float32, device=device) - if not is_timesteps_provided: - timesteps = sigmas * self.config.num_train_timesteps - else: - timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32, device=device) + timesteps = sigmas * self.config.num_train_timesteps # 6. Append the terminal sigma value. # If a model requires inverted sigma schedule for denoising but timesteps without inversion, the diff --git a/src/diffusers/schedulers/scheduling_lcm.py b/src/diffusers/schedulers/scheduling_lcm.py index e70bc4745a24..9c47fefac935 100644 --- a/src/diffusers/schedulers/scheduling_lcm.py +++ b/src/diffusers/schedulers/scheduling_lcm.py @@ -40,9 +40,9 @@ class LCMSchedulerOutput(BaseOutput): prev_sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images): Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the denoising loop. - pred_original_sample (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images): - The predicted denoised sample `(x_{0})` based on the model output from the current timestep. - `pred_original_sample` can be used to preview progress or for guidance. + denoised (`torch.Tensor` of shape `(batch_size, num_channels, height, width)` for images, *optional*): + The predicted denoised sample `(x_{0})` based on the model output from the current timestep. `denoised` can + be used to preview progress or for guidance. """ prev_sample: torch.Tensor @@ -150,45 +150,45 @@ class LCMScheduler(SchedulerMixin, ConfigMixin): functionality via the [`SchedulerMixin.save_pretrained`] and [`~SchedulerMixin.from_pretrained`] functions. Args: - num_train_timesteps (`int`, defaults to 1000): + num_train_timesteps (`int`, defaults to `1000`): The number of diffusion steps to train the model. - beta_start (`float`, defaults to 0.0001): + beta_start (`float`, defaults to `0.00085`): The starting `beta` value of inference. - beta_end (`float`, defaults to 0.02): + beta_end (`float`, defaults to `0.012`): The final `beta` value. - beta_schedule (`str`, defaults to `"linear"`): + beta_schedule (`str`, defaults to `"scaled_linear"`): The beta schedule, a mapping from a beta range to a sequence of betas for stepping the model. Choose from `linear`, `scaled_linear`, or `squaredcos_cap_v2`. - trained_betas (`np.ndarray`, *optional*): + trained_betas (`np.ndarray` or `list[float]`, *optional*): Pass an array of betas directly to the constructor to bypass `beta_start` and `beta_end`. - original_inference_steps (`int`, *optional*, defaults to 50): + original_inference_steps (`int`, defaults to `50`): The default number of inference steps used to generate a linearly-spaced timestep schedule, from which we will ultimately take `num_inference_steps` evenly spaced timesteps to form the final timestep schedule. - clip_sample (`bool`, defaults to `True`): + clip_sample (`bool`, defaults to `False`): Clip the predicted sample for numerical stability. - clip_sample_range (`float`, defaults to 1.0): + clip_sample_range (`float`, defaults to `1.0`): The maximum magnitude for sample clipping. Valid only when `clip_sample=True`. set_alpha_to_one (`bool`, defaults to `True`): Each diffusion step uses the alphas product value at that step and at the previous one. For the final step there is no previous alpha. When this option is `True` the previous alpha product is fixed to `1`, otherwise it uses the alpha value at step 0. - steps_offset (`int`, defaults to 0): + steps_offset (`int`, defaults to `0`): An offset added to the inference steps, as required by some model families. - prediction_type (`str`, defaults to `epsilon`, *optional*): + prediction_type (`str`, defaults to `"epsilon"`): Prediction type of the scheduler function; can be `epsilon` (predicts the noise of the diffusion process), `sample` (directly predicts the noisy sample`) or `v_prediction` (see section 2.4 of [Imagen Video](https://huggingface.co/papers/2210.02303) paper). thresholding (`bool`, defaults to `False`): Whether to use the "dynamic thresholding" method. This is unsuitable for latent-space diffusion models such as Stable Diffusion. - dynamic_thresholding_ratio (`float`, defaults to 0.995): + dynamic_thresholding_ratio (`float`, defaults to `0.995`): The ratio for the dynamic thresholding method. Valid only when `thresholding=True`. - sample_max_value (`float`, defaults to 1.0): + sample_max_value (`float`, defaults to `1.0`): The threshold value for dynamic thresholding. Valid only when `thresholding=True`. timestep_spacing (`str`, defaults to `"leading"`): The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information. - timestep_scaling (`float`, defaults to 10.0): + timestep_scaling (`float`, defaults to `10.0`): The factor the timesteps will be multiplied by when calculating the consistency model boundary conditions `c_skip` and `c_out`. Increasing this will decrease the approximation error (although the approximation error at the default of `10.0` is already pretty small). @@ -306,13 +306,24 @@ def _init_step_index(self, timestep: float | torch.Tensor) -> None: self._step_index = self._begin_index @property - def step_index(self): + def step_index(self) -> int | None: + """ + The index counter for current timestep. It will increase by 1 after each scheduler step. + + Returns: + `int` or `None`: + The current step index, or `None` if not initialized. + """ return self._step_index @property - def begin_index(self): + def begin_index(self) -> int | None: """ The index for the first timestep. It should be set from pipeline with `set_begin_index` method. + + Returns: + `int` or `None`: + The begin index for the scheduler, or `None` if not set. """ return self._begin_index @@ -337,6 +348,7 @@ def scale_model_input(self, sample: torch.Tensor, timestep: int | None = None) - The input sample. timestep (`int`, *optional*): The current timestep in the diffusion chain. + Returns: `torch.Tensor`: A scaled input sample. @@ -390,11 +402,11 @@ def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor: def set_timesteps( self, num_inference_steps: int | None = None, - device: str | torch.device = None, + device: str | torch.device | None = None, original_inference_steps: int | None = None, timesteps: list[int] | None = None, - strength: int = 1.0, - ): + strength: float = 1.0, + ) -> None: """ Sets the discrete timesteps used for the diffusion chain (to be run before inference). @@ -413,6 +425,12 @@ def set_timesteps( Custom timesteps used to support arbitrary spacing between timesteps. If `None`, then the default timestep spacing strategy of equal spacing between timesteps on the training/distillation timestep schedule is used. If `timesteps` is passed, `num_inference_steps` must be `None`. + strength (`float`, defaults to `1.0`): + Strength factor used to generate a partial timestep schedule (e.g. for image-to-image). A value of + `1.0` uses the full schedule. + + Returns: + `None` """ # 0. Check inputs if num_inference_steps is None and timesteps is None: @@ -527,7 +545,20 @@ def set_timesteps( self._step_index = None self._begin_index = None - def get_scalings_for_boundary_condition_discrete(self, timestep): + def get_scalings_for_boundary_condition_discrete(self, timestep: int) -> tuple[float, float]: + """ + Computes the boundary condition scalings (`c_skip` and `c_out`) for the given discrete timestep, as used in the + Latent Consistency Model. + + Args: + timestep (`int`): + The discrete timestep for which to compute the scalings. + + Returns: + `tuple[float, float]`: + A tuple containing `c_skip` (scaling for the input sample) and `c_out` (scaling for the predicted + denoised sample). + """ self.sigma_data = 0.5 # Default: 0.5 scaled_timestep = timestep * self.config.timestep_scaling @@ -550,16 +581,17 @@ def step( Args: model_output (`torch.Tensor`): The direct output from learned diffusion model. - timestep (`float`): + timestep (`int`): The current discrete timestep in the diffusion chain. sample (`torch.Tensor`): A current instance of a sample created by the diffusion process. generator (`torch.Generator`, *optional*): - A random number generator. - return_dict (`bool`, *optional*, defaults to `True`): + A random number generator for reproducible sampling. + return_dict (`bool`, defaults to `True`): Whether or not to return a [`~schedulers.scheduling_lcm.LCMSchedulerOutput`] or `tuple`. + Returns: - [`~schedulers.scheduling_utils.LCMSchedulerOutput`] or `tuple`: + [`~schedulers.scheduling_lcm.LCMSchedulerOutput`] or `tuple`: If return_dict is `True`, [`~schedulers.scheduling_lcm.LCMSchedulerOutput`] is returned, otherwise a tuple is returned where the first element is the sample tensor. """ @@ -709,7 +741,14 @@ def get_velocity(self, sample: torch.Tensor, noise: torch.Tensor, timesteps: tor velocity = sqrt_alpha_prod * noise - sqrt_one_minus_alpha_prod * sample return velocity - def __len__(self): + def __len__(self) -> int: + """ + Returns the number of train timesteps. + + Returns: + `int`: + The number of diffusion steps used to train the model. + """ return self.config.num_train_timesteps # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler.previous_timestep diff --git a/src/diffusers/schedulers/scheduling_scm.py b/src/diffusers/schedulers/scheduling_scm.py index fb3bbc40d726..d3ca01dee507 100644 --- a/src/diffusers/schedulers/scheduling_scm.py +++ b/src/diffusers/schedulers/scheduling_scm.py @@ -55,11 +55,11 @@ class SCMScheduler(SchedulerMixin, ConfigMixin): documentation for the generic methods the library implements for all schedulers such as loading and saving. Args: - num_train_timesteps (`int`, defaults to 1000): + num_train_timesteps (`int`, defaults to `1000`): The number of diffusion steps to train the model. prediction_type (`str`, defaults to `trigflow`): Prediction type of the scheduler function. Currently only supports "trigflow". - sigma_data (`float`, defaults to 0.5): + sigma_data (`float`, defaults to `0.5`): The standard deviation of the noise added during multi-step inference. """ @@ -77,11 +77,11 @@ def __init__( Initialize the SCM scheduler. Args: - num_train_timesteps (`int`, defaults to 1000): + num_train_timesteps (`int`, defaults to `1000`): The number of diffusion steps to train the model. prediction_type (`str`, defaults to `trigflow`): Prediction type of the scheduler function. Currently only supports "trigflow". - sigma_data (`float`, defaults to 0.5): + sigma_data (`float`, defaults to `0.5`): The standard deviation of the noise added during multi-step inference. """ # standard deviation of the initial noise distribution @@ -95,11 +95,25 @@ def __init__( self._begin_index = None @property - def step_index(self): + def step_index(self) -> int: + """ + The index counter for current timestep. It will increase by 1 after each scheduler step. + + Returns: + `int` or `None`: + The current step index, or `None` if not initialized. + """ return self._step_index @property - def begin_index(self): + def begin_index(self) -> int: + """ + The index for the first timestep. It should be set from pipeline with `set_begin_index` method. + + Returns: + `int` or `None`: + The begin index for the scheduler, or `None` if not set. + """ return self._begin_index # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index @@ -116,10 +130,10 @@ def set_begin_index(self, begin_index: int = 0): def set_timesteps( self, num_inference_steps: int, - timesteps: torch.Tensor = None, - device: str | torch.device = None, - max_timesteps: float = 1.57080, - intermediate_timesteps: float = 1.3, + timesteps: list[float] | torch.Tensor | None = None, + device: str | torch.device | None = None, + max_timesteps: float | None = 1.57080, + intermediate_timesteps: float | None = 1.3, ): """ Sets the discrete timesteps used for the diffusion chain (to be run before inference). @@ -127,11 +141,13 @@ def set_timesteps( Args: num_inference_steps (`int`): The number of diffusion steps used when generating samples with a pre-trained model. - timesteps (`torch.Tensor`, *optional*): + timesteps (`list[float]` or `torch.Tensor`, *optional*): Custom timesteps to use for the denoising process. - max_timesteps (`float`, defaults to 1.57080): + device (`str` or `torch.device`, *optional*): + The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. + max_timesteps (`float`, *optional*, defaults to `1.57080`): The maximum timestep value used in the SCM scheduler. - intermediate_timesteps (`float`, *optional*, defaults to 1.3): + intermediate_timesteps (`float`, *optional*, defaults to `1.3`): The intermediate timestep value used in SCM scheduler (only used when num_inference_steps=2). """ if num_inference_steps > self.config.num_train_timesteps: @@ -223,7 +239,7 @@ def step( model_output: torch.FloatTensor, timestep: float, sample: torch.FloatTensor, - generator: torch.Generator = None, + generator: torch.Generator | None = None, return_dict: bool = True, ) -> SCMSchedulerOutput | tuple: """ @@ -237,10 +253,13 @@ def step( The current discrete timestep in the diffusion chain. sample (`torch.FloatTensor`): A current instance of a sample created by the diffusion process. + generator (`torch.Generator`, *optional*): + A random number generator for reproducible sampling. return_dict (`bool`, *optional*, defaults to `True`): Whether or not to return a [`~schedulers.scheduling_scm.SCMSchedulerOutput`] or `tuple`. + Returns: - [`~schedulers.scheduling_utils.SCMSchedulerOutput`] or `tuple`: + [`~schedulers.scheduling_scm.SCMSchedulerOutput`] or `tuple`: If return_dict is `True`, [`~schedulers.scheduling_scm.SCMSchedulerOutput`] is returned, otherwise a tuple is returned where the first element is the sample tensor. """ @@ -282,5 +301,5 @@ def step( return SCMSchedulerOutput(prev_sample=prev_sample, pred_original_sample=pred_x0) - def __len__(self): + def __len__(self) -> int: return self.config.num_train_timesteps diff --git a/src/diffusers/utils/dummy_pt_objects.py b/src/diffusers/utils/dummy_pt_objects.py index b85e9b1a1d39..9035efb3e6e2 100644 --- a/src/diffusers/utils/dummy_pt_objects.py +++ b/src/diffusers/utils/dummy_pt_objects.py @@ -1485,6 +1485,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class JoyImageEditPlusTransformer3DModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class JoyImageEditTransformer3DModel(metaclass=DummyObject): _backends = ["torch"] diff --git a/src/diffusers/utils/dummy_torch_and_transformers_objects.py b/src/diffusers/utils/dummy_torch_and_transformers_objects.py index ed4e13a57eb1..b03bc6926197 100644 --- a/src/diffusers/utils/dummy_torch_and_transformers_objects.py +++ b/src/diffusers/utils/dummy_torch_and_transformers_objects.py @@ -32,6 +32,36 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch", "transformers"]) +class Cosmos3OmniBlocks(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + +class Cosmos3OmniModularPipeline(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + class ErnieImageAutoBlocks(metaclass=DummyObject): _backends = ["torch", "transformers"] @@ -2252,6 +2282,36 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch", "transformers"]) +class JoyImageEditPlusPipeline(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + +class JoyImageEditPlusPipelineOutput(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + class Kandinsky3Img2ImgPipeline(metaclass=DummyObject): _backends = ["torch", "transformers"] diff --git a/src/diffusers/utils/remote_utils.py b/src/diffusers/utils/remote_utils.py index d109e697dd11..3a478ecf4085 100644 --- a/src/diffusers/utils/remote_utils.py +++ b/src/diffusers/utils/remote_utils.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/hooks/test_group_offloading.py b/tests/hooks/test_group_offloading.py index 7e7741ed77d3..a903186aa6b4 100644 --- a/tests/hooks/test_group_offloading.py +++ b/tests/hooks/test_group_offloading.py @@ -1,4 +1,4 @@ -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/hooks/test_hooks.py b/tests/hooks/test_hooks.py index 26418adfddee..b40d9478355d 100644 --- a/tests/hooks/test_hooks.py +++ b/tests/hooks/test_hooks.py @@ -1,4 +1,4 @@ -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/hooks/test_mag_cache.py b/tests/hooks/test_mag_cache.py index e5a500d8b291..4fd45a73e326 100644 --- a/tests/hooks/test_mag_cache.py +++ b/tests/hooks/test_mag_cache.py @@ -1,4 +1,4 @@ -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/lora/test_lora_layers_auraflow.py b/tests/lora/test_lora_layers_auraflow.py index 78ef4ce151be..ba55cd7baac3 100644 --- a/tests/lora/test_lora_layers_auraflow.py +++ b/tests/lora/test_lora_layers_auraflow.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/lora/test_lora_layers_cogvideox.py b/tests/lora/test_lora_layers_cogvideox.py index 7bd54b77ca35..33fc8f39046e 100644 --- a/tests/lora/test_lora_layers_cogvideox.py +++ b/tests/lora/test_lora_layers_cogvideox.py @@ -1,4 +1,4 @@ -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/lora/test_lora_layers_cogview4.py b/tests/lora/test_lora_layers_cogview4.py index e8ee6e7a7db6..dfc387369f52 100644 --- a/tests/lora/test_lora_layers_cogview4.py +++ b/tests/lora/test_lora_layers_cogview4.py @@ -1,4 +1,4 @@ -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/lora/test_lora_layers_flux.py b/tests/lora/test_lora_layers_flux.py index b840d7ac72ce..4db03af3e292 100644 --- a/tests/lora/test_lora_layers_flux.py +++ b/tests/lora/test_lora_layers_flux.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/lora/test_lora_layers_flux2.py b/tests/lora/test_lora_layers_flux2.py index d970b7d7847f..25451fa51957 100644 --- a/tests/lora/test_lora_layers_flux2.py +++ b/tests/lora/test_lora_layers_flux2.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/lora/test_lora_layers_helios.py b/tests/lora/test_lora_layers_helios.py index fbcc3b808eee..44dae796070c 100644 --- a/tests/lora/test_lora_layers_helios.py +++ b/tests/lora/test_lora_layers_helios.py @@ -1,4 +1,4 @@ -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/lora/test_lora_layers_hunyuanvideo.py b/tests/lora/test_lora_layers_hunyuanvideo.py index e59bc5662fe1..06f9813ea906 100644 --- a/tests/lora/test_lora_layers_hunyuanvideo.py +++ b/tests/lora/test_lora_layers_hunyuanvideo.py @@ -1,4 +1,4 @@ -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/lora/test_lora_layers_ltx2.py b/tests/lora/test_lora_layers_ltx2.py index 0a4b14454f5b..f75796f75b0e 100644 --- a/tests/lora/test_lora_layers_ltx2.py +++ b/tests/lora/test_lora_layers_ltx2.py @@ -1,4 +1,4 @@ -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/lora/test_lora_layers_ltx_video.py b/tests/lora/test_lora_layers_ltx_video.py index 095e5b577cf0..6fbae88281ce 100644 --- a/tests/lora/test_lora_layers_ltx_video.py +++ b/tests/lora/test_lora_layers_ltx_video.py @@ -1,4 +1,4 @@ -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/lora/test_lora_layers_lumina2.py b/tests/lora/test_lora_layers_lumina2.py index da032229a785..7dfa7d8956aa 100644 --- a/tests/lora/test_lora_layers_lumina2.py +++ b/tests/lora/test_lora_layers_lumina2.py @@ -1,4 +1,4 @@ -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/lora/test_lora_layers_mochi.py b/tests/lora/test_lora_layers_mochi.py index ee8254112924..da6543900f5d 100644 --- a/tests/lora/test_lora_layers_mochi.py +++ b/tests/lora/test_lora_layers_mochi.py @@ -1,4 +1,4 @@ -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/lora/test_lora_layers_qwenimage.py b/tests/lora/test_lora_layers_qwenimage.py index 73fd026a670c..480698af58ea 100644 --- a/tests/lora/test_lora_layers_qwenimage.py +++ b/tests/lora/test_lora_layers_qwenimage.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/lora/test_lora_layers_sana.py b/tests/lora/test_lora_layers_sana.py index 97bf5cbba920..8406669710e4 100644 --- a/tests/lora/test_lora_layers_sana.py +++ b/tests/lora/test_lora_layers_sana.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/lora/test_lora_layers_sd.py b/tests/lora/test_lora_layers_sd.py index 933bf2336a59..4662075cfd0a 100644 --- a/tests/lora/test_lora_layers_sd.py +++ b/tests/lora/test_lora_layers_sd.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/lora/test_lora_layers_sd3.py b/tests/lora/test_lora_layers_sd3.py index 228460eaad90..11db0823ed11 100644 --- a/tests/lora/test_lora_layers_sd3.py +++ b/tests/lora/test_lora_layers_sd3.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/lora/test_lora_layers_sdxl.py b/tests/lora/test_lora_layers_sdxl.py index ac1d65abdaa7..5ed9fa53a346 100644 --- a/tests/lora/test_lora_layers_sdxl.py +++ b/tests/lora/test_lora_layers_sdxl.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/lora/test_lora_layers_wan.py b/tests/lora/test_lora_layers_wan.py index 1a5b47ce4e5e..b0fb000dd2c2 100644 --- a/tests/lora/test_lora_layers_wan.py +++ b/tests/lora/test_lora_layers_wan.py @@ -1,4 +1,4 @@ -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/lora/test_lora_layers_wanvace.py b/tests/lora/test_lora_layers_wanvace.py index dc435094d780..4089acd1a8d1 100644 --- a/tests/lora/test_lora_layers_wanvace.py +++ b/tests/lora/test_lora_layers_wanvace.py @@ -1,4 +1,4 @@ -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/lora/test_lora_layers_z_image.py b/tests/lora/test_lora_layers_z_image.py index 8432ea56a6fa..5bdaf9147667 100644 --- a/tests/lora/test_lora_layers_z_image.py +++ b/tests/lora/test_lora_layers_z_image.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/lora/utils.py b/tests/lora/utils.py index a3b0fe0aa2f2..38aec8ce4807 100644 --- a/tests/lora/utils.py +++ b/tests/lora/utils.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/autoencoders/test_models_asymmetric_autoencoder_kl.py b/tests/models/autoencoders/test_models_asymmetric_autoencoder_kl.py index 149ca395a328..425848efd700 100644 --- a/tests/models/autoencoders/test_models_asymmetric_autoencoder_kl.py +++ b/tests/models/autoencoders/test_models_asymmetric_autoencoder_kl.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/autoencoders/test_models_autoencoder_cosmos.py b/tests/models/autoencoders/test_models_autoencoder_cosmos.py index 2aa8fbec380e..b70f05dc47eb 100644 --- a/tests/models/autoencoders/test_models_autoencoder_cosmos.py +++ b/tests/models/autoencoders/test_models_autoencoder_cosmos.py @@ -1,4 +1,4 @@ -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/autoencoders/test_models_autoencoder_dc.py b/tests/models/autoencoders/test_models_autoencoder_dc.py index f0ed1816ce70..a5ce0a975a11 100644 --- a/tests/models/autoencoders/test_models_autoencoder_dc.py +++ b/tests/models/autoencoders/test_models_autoencoder_dc.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/autoencoders/test_models_autoencoder_hunyuan_video.py b/tests/models/autoencoders/test_models_autoencoder_hunyuan_video.py index 35d4244e0b81..cdaebcf2eea9 100644 --- a/tests/models/autoencoders/test_models_autoencoder_hunyuan_video.py +++ b/tests/models/autoencoders/test_models_autoencoder_hunyuan_video.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/autoencoders/test_models_autoencoder_kl.py b/tests/models/autoencoders/test_models_autoencoder_kl.py index 0e2297f22e4c..0e3e1aa3a46a 100644 --- a/tests/models/autoencoders/test_models_autoencoder_kl.py +++ b/tests/models/autoencoders/test_models_autoencoder_kl.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/autoencoders/test_models_autoencoder_kl_cogvideox.py b/tests/models/autoencoders/test_models_autoencoder_kl_cogvideox.py index cfb1420196a7..169eca51c65a 100644 --- a/tests/models/autoencoders/test_models_autoencoder_kl_cogvideox.py +++ b/tests/models/autoencoders/test_models_autoencoder_kl_cogvideox.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/autoencoders/test_models_autoencoder_kl_kvae.py b/tests/models/autoencoders/test_models_autoencoder_kl_kvae.py index 5256acee442b..066864abb4ff 100644 --- a/tests/models/autoencoders/test_models_autoencoder_kl_kvae.py +++ b/tests/models/autoencoders/test_models_autoencoder_kl_kvae.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/autoencoders/test_models_autoencoder_kl_kvae_video.py b/tests/models/autoencoders/test_models_autoencoder_kl_kvae_video.py index dc74842db8b8..a7097e128a13 100644 --- a/tests/models/autoencoders/test_models_autoencoder_kl_kvae_video.py +++ b/tests/models/autoencoders/test_models_autoencoder_kl_kvae_video.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/autoencoders/test_models_autoencoder_kl_ltx2_audio.py b/tests/models/autoencoders/test_models_autoencoder_kl_ltx2_audio.py index 2e16ba3f9953..a3a1379a2f2f 100644 --- a/tests/models/autoencoders/test_models_autoencoder_kl_ltx2_audio.py +++ b/tests/models/autoencoders/test_models_autoencoder_kl_ltx2_audio.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/autoencoders/test_models_autoencoder_kl_temporal_decoder.py b/tests/models/autoencoders/test_models_autoencoder_kl_temporal_decoder.py index 3958fccae936..56297b7540bb 100644 --- a/tests/models/autoencoders/test_models_autoencoder_kl_temporal_decoder.py +++ b/tests/models/autoencoders/test_models_autoencoder_kl_temporal_decoder.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/autoencoders/test_models_autoencoder_ltx2_video.py b/tests/models/autoencoders/test_models_autoencoder_ltx2_video.py index cc041baa5bc7..dac3f8fdcc6b 100644 --- a/tests/models/autoencoders/test_models_autoencoder_ltx2_video.py +++ b/tests/models/autoencoders/test_models_autoencoder_ltx2_video.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/autoencoders/test_models_autoencoder_ltx_video.py b/tests/models/autoencoders/test_models_autoencoder_ltx_video.py index 2269f1a0136b..d93c1e566712 100644 --- a/tests/models/autoencoders/test_models_autoencoder_ltx_video.py +++ b/tests/models/autoencoders/test_models_autoencoder_ltx_video.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/autoencoders/test_models_autoencoder_magvit.py b/tests/models/autoencoders/test_models_autoencoder_magvit.py index f7c484dbbed2..2f3ad90d40e1 100644 --- a/tests/models/autoencoders/test_models_autoencoder_magvit.py +++ b/tests/models/autoencoders/test_models_autoencoder_magvit.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/autoencoders/test_models_autoencoder_mochi.py b/tests/models/autoencoders/test_models_autoencoder_mochi.py index 71fd76905451..3beb0c94066b 100755 --- a/tests/models/autoencoders/test_models_autoencoder_mochi.py +++ b/tests/models/autoencoders/test_models_autoencoder_mochi.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/autoencoders/test_models_autoencoder_oobleck.py b/tests/models/autoencoders/test_models_autoencoder_oobleck.py index ccd67bc7c704..bf603da829ce 100644 --- a/tests/models/autoencoders/test_models_autoencoder_oobleck.py +++ b/tests/models/autoencoders/test_models_autoencoder_oobleck.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/autoencoders/test_models_autoencoder_rae.py b/tests/models/autoencoders/test_models_autoencoder_rae.py index cc8869737bcc..c2b801ec348c 100644 --- a/tests/models/autoencoders/test_models_autoencoder_rae.py +++ b/tests/models/autoencoders/test_models_autoencoder_rae.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/autoencoders/test_models_autoencoder_tiny.py b/tests/models/autoencoders/test_models_autoencoder_tiny.py index 43dda6187505..8f0740a8a70e 100644 --- a/tests/models/autoencoders/test_models_autoencoder_tiny.py +++ b/tests/models/autoencoders/test_models_autoencoder_tiny.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/autoencoders/test_models_autoencoder_vidtok.py b/tests/models/autoencoders/test_models_autoencoder_vidtok.py index 087dca5debfa..9810296a07d9 100644 --- a/tests/models/autoencoders/test_models_autoencoder_vidtok.py +++ b/tests/models/autoencoders/test_models_autoencoder_vidtok.py @@ -13,7 +13,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import pytest import torch from diffusers import AutoencoderVidTok @@ -80,7 +79,6 @@ def get_dummy_inputs(self) -> dict: class TestAutoencoderVidTok(AutoencoderVidTokTesterConfig, ModelTesterMixin): - @pytest.mark.skip("VidTok output structure not compatible with recursive output check.") def test_outputs_equivalence(self): super().test_outputs_equivalence() diff --git a/tests/models/autoencoders/test_models_autoencoder_wan.py b/tests/models/autoencoders/test_models_autoencoder_wan.py index b5b89769e97e..89e5c58e4a27 100644 --- a/tests/models/autoencoders/test_models_autoencoder_wan.py +++ b/tests/models/autoencoders/test_models_autoencoder_wan.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/autoencoders/test_models_consistency_decoder_vae.py b/tests/models/autoencoders/test_models_consistency_decoder_vae.py index 906baa60a9dc..b5831d6ba09d 100644 --- a/tests/models/autoencoders/test_models_consistency_decoder_vae.py +++ b/tests/models/autoencoders/test_models_consistency_decoder_vae.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -16,6 +16,7 @@ import gc import numpy as np +import pytest import torch from diffusers import ConsistencyDecoderVAE, StableDiffusionPipeline @@ -86,7 +87,12 @@ def get_dummy_inputs(self) -> dict: class TestConsistencyDecoderVAE(ConsistencyDecoderVAETesterConfig, ModelTesterMixin): - pass + @pytest.mark.skip( + reason="The consistency decoder samples noise (`randn_tensor`) during `decode`, so two forward passes " + "diverge regardless of dtype. This makes a save/load output comparison non-deterministic." + ) + def test_from_save_pretrained_dtype_inference(self, *args, **kwargs): + pass class TestConsistencyDecoderVAETraining(ConsistencyDecoderVAETesterConfig, TrainingTesterMixin): @@ -212,7 +218,7 @@ def test_vae_tiling(self): generator=torch.Generator("cpu").manual_seed(0), ).images[0] - pipe.enable_vae_tiling() + pipe.vae.enable_tiling() out_2 = pipe( "horse", num_inference_steps=2, diff --git a/tests/models/autoencoders/test_models_vq.py b/tests/models/autoencoders/test_models_vq.py index ce1606f0e859..b6dbcfa64b19 100644 --- a/tests/models/autoencoders/test_models_vq.py +++ b/tests/models/autoencoders/test_models_vq.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -64,6 +64,17 @@ def get_dummy_inputs(self) -> dict: class TestVQModel(VQModelTesterConfig, ModelTesterMixin): + @pytest.mark.skipif( + torch_device not in ["cuda", "xpu"], + reason="float16 and bfloat16 can only be use for inference with an accelerator", + ) + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16], ids=["fp16", "bf16"]) + def test_from_save_pretrained_dtype_inference(self, tmp_path, dtype): + # The reference and reloaded models hold identical weights, so any output difference is + # half-precision kernel nondeterminism between the two module instances rather than a save/load + # fidelity issue. The default 1e-4 tolerance is too tight for that fp16/bf16 noise on some GPUs. + super().test_from_save_pretrained_dtype_inference(tmp_path, dtype, atol=1e-3) + def test_from_pretrained_hub(self): model, loading_info = VQModel.from_pretrained("fusing/vqgan-dummy", output_loading_info=True) assert model is not None diff --git a/tests/models/controlnets/test_models_controlnet_cosmos.py b/tests/models/controlnets/test_models_controlnet_cosmos.py index e7ea6362213d..f6a3b35006f6 100644 --- a/tests/models/controlnets/test_models_controlnet_cosmos.py +++ b/tests/models/controlnets/test_models_controlnet_cosmos.py @@ -1,4 +1,4 @@ -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -251,6 +251,10 @@ def test_determinism(self): def test_from_save_pretrained(self): super().test_from_save_pretrained() + @pytest.mark.skip("Output is a list of tensors; comparison helper calls .shape on it.") + def test_from_save_pretrained_dtype_inference(self, *args, **kwargs): + super().test_from_save_pretrained_dtype_inference(*args, **kwargs) + @pytest.mark.skip("Output is a list of tensors; comparison helper calls .shape on it.") def test_from_save_pretrained_variant(self): super().test_from_save_pretrained_variant() diff --git a/tests/models/test_layers_utils.py b/tests/models/test_layers_utils.py index eaeffa699db2..bcdf64abe6b8 100644 --- a/tests/models/test_layers_utils.py +++ b/tests/models/test_layers_utils.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/test_modeling_common.py b/tests/models/test_modeling_common.py index 7e7822ac16ea..0022e2c1c833 100644 --- a/tests/models/test_modeling_common.py +++ b/tests/models/test_modeling_common.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/testing_utils/attention.py b/tests/models/testing_utils/attention.py index 8672e19e6528..f31323d1bf52 100644 --- a/tests/models/testing_utils/attention.py +++ b/tests/models/testing_utils/attention.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -70,8 +70,38 @@ ], ) +_PARAM_FLASH_VARLEN_HUB = pytest.param( + AttentionBackendName.FLASH_VARLEN_HUB, + id="flash_varlen_hub", + marks=[ + pytest.mark.skipif(not _CUDA_AVAILABLE, reason="CUDA is required for flash_varlen_hub backend."), + pytest.mark.skipif( + not is_kernels_available(), + reason="`kernels` package is required for flash_varlen_hub backend. Install with `pip install kernels`.", + ), + ], +) + +_PARAM_FLASH_3_VARLEN_HUB = pytest.param( + AttentionBackendName._FLASH_3_VARLEN_HUB, + id="flash_3_varlen_hub", + marks=[ + pytest.mark.skipif(not _CUDA_AVAILABLE, reason="CUDA is required for _flash_3_varlen_hub backend."), + pytest.mark.skipif( + not is_kernels_available(), + reason="`kernels` package is required for _flash_3_varlen_hub backend. Install with `pip install kernels`.", + ), + ], +) + # All backends under test. -_ALL_BACKEND_PARAMS = [_PARAM_NATIVE_CUDNN, _PARAM_FLASH_HUB, _PARAM_FLASH_3_HUB] +_ALL_BACKEND_PARAMS = [ + _PARAM_NATIVE_CUDNN, + _PARAM_FLASH_HUB, + _PARAM_FLASH_3_HUB, + _PARAM_FLASH_VARLEN_HUB, + _PARAM_FLASH_3_VARLEN_HUB, +] # Backends that perform non-deterministic operations and therefore cannot run when # torch.use_deterministic_algorithms(True) is active (e.g. after enable_full_determinism()). @@ -267,6 +297,8 @@ class AttentionBackendTesterMixin: Use `pytest -m "not attention"` to skip these tests. """ + unsupported_attn_backends: list[str] = [] + def setup_method(self): gc.collect() backend_empty_cache(torch_device) @@ -280,6 +312,8 @@ def teardown_method(self): def test_set_attention_backend_matches_context_manager(self, backend): """set_attention_backend() and the attention_backend() context manager must yield identical outputs.""" _skip_if_backend_requires_nondeterminism(backend) + if backend.value in self.unsupported_attn_backends: + pytest.skip(f"{backend.value} is not supported for this model.") init_dict = self.get_init_dict() inputs_dict = self.get_dummy_inputs() @@ -318,6 +352,8 @@ def test_set_attention_backend_matches_context_manager(self, backend): def test_output_close_to_native(self, backend, atol=1e-2, rtol=1e-2): """All backends should produce model output numerically close to the native SDPA reference.""" _skip_if_backend_requires_nondeterminism(backend) + if backend.value in self.unsupported_attn_backends: + pytest.skip(f"{backend.value} is not supported for this model.") init_dict = self.get_init_dict() inputs_dict = self.get_dummy_inputs() @@ -378,6 +414,8 @@ def test_compile(self, backend, atol=1e-2, rtol=1e-2): as opposed to `model.compile`). """ _skip_if_backend_requires_nondeterminism(backend) + if backend.value in self.unsupported_attn_backends: + pytest.skip(f"{backend.value} is not supported for this model.") if getattr(self.model_class, "_repeated_blocks", None) is None: pytest.skip("Skipping tests as regional compilation is not supported.") diff --git a/tests/models/testing_utils/cache.py b/tests/models/testing_utils/cache.py index e8a835f6bccf..8357d34f3077 100644 --- a/tests/models/testing_utils/cache.py +++ b/tests/models/testing_utils/cache.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/testing_utils/common.py b/tests/models/testing_utils/common.py index eb120567f3d1..570957cdaf74 100644 --- a/tests/models/testing_utils/common.py +++ b/tests/models/testing_utils/common.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -507,6 +507,29 @@ def test_keep_in_fp32_modules(self, tmp_path): else: assert param.dtype == torch.float16, f"Parameter {name} should be float16 but got {param.dtype}" + def test_to_keep_in_fp32_modules_warns(self, caplog): + fp32_modules = self.model_class._keep_in_fp32_modules + if fp32_modules is None or len(fp32_modules) == 0: + pytest.skip("Model does not have _keep_in_fp32_modules defined.") + + model = self.model_class(**self.get_init_dict()) + + logger_name = "diffusers.models.modeling_utils" + logging.enable_propagation() + try: + with caplog.at_level(logging.WARNING, logger=logger_name): + caplog.clear() + model.to(torch.float16) + finally: + logging.disable_propagation() + + expected_message = ( + f"There are modules in {model.__class__.__name__} that should be kept in float32: " + f"{fp32_modules}. Casting directly with `to()` can lead to inconsistent results; set " + f"`torch_dtype` in `from_pretrained()` instead to keep these modules in float32." + ) + assert expected_message in caplog.text + @require_accelerator @pytest.mark.skipif( torch_device not in ["cuda", "xpu"], @@ -519,7 +542,25 @@ def test_from_save_pretrained_dtype_inference(self, tmp_path, dtype, atol=1e-4, model.to(torch_device) fp32_modules = model._keep_in_fp32_modules or [] - model.to(dtype).save_pretrained(tmp_path) + # Build the reference model with the same mixed-precision layout that `from_pretrained` enforces, so + # the comparison reflects real save/load fidelity: + # - `_keep_in_fp32_modules` stay in fp32 while everything else is cast to `dtype`; + # - non-persistent buffers (e.g. fp32 RoPE `inv_freq`) are left untouched, because they are not part + # of the checkpoint and are regenerated by `__init__` on load. Truncating them here would make the + # reference diverge from the reloaded model for reasons unrelated to save/load. + persistent_tensor_names = {name for name, _ in named_persistent_module_tensors(model, recurse=True)} + + def keep_in_fp32(name): + return any(module_to_keep_in_fp32 in name.split(".") for module_to_keep_in_fp32 in fp32_modules) + + for name, param in model.named_parameters(): + param.data = param.data.to(torch.float32 if keep_in_fp32(name) else dtype) + for name, buf in model.named_buffers(): + if not buf.is_floating_point() or name not in persistent_tensor_names: + continue + buf.data = buf.data.to(torch.float32 if keep_in_fp32(name) else dtype) + + model.save_pretrained(tmp_path) model_loaded = self.model_class.from_pretrained(tmp_path, torch_dtype=dtype).to(torch_device) for name, param in model_loaded.named_parameters(): diff --git a/tests/models/testing_utils/compile.py b/tests/models/testing_utils/compile.py index 4787d0742b18..5d638ec0b6ff 100644 --- a/tests/models/testing_utils/compile.py +++ b/tests/models/testing_utils/compile.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/testing_utils/ip_adapter.py b/tests/models/testing_utils/ip_adapter.py index 632019c87499..b596469e7ae7 100644 --- a/tests/models/testing_utils/ip_adapter.py +++ b/tests/models/testing_utils/ip_adapter.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/testing_utils/lora.py b/tests/models/testing_utils/lora.py index dfa326b014d0..3fccc5134155 100644 --- a/tests/models/testing_utils/lora.py +++ b/tests/models/testing_utils/lora.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/testing_utils/memory.py b/tests/models/testing_utils/memory.py index 9cb28919efb8..a2b8e06a3dc7 100644 --- a/tests/models/testing_utils/memory.py +++ b/tests/models/testing_utils/memory.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/testing_utils/parallelism.py b/tests/models/testing_utils/parallelism.py index d6b8854f93bf..dad2f51f1f2c 100644 --- a/tests/models/testing_utils/parallelism.py +++ b/tests/models/testing_utils/parallelism.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -426,6 +426,10 @@ class ContextParallelAttentionBackendsTesterMixin: "_flash_3_hub", marks=pytest.mark.skipif(not is_kernels_available(), reason="`kernels` is not available."), ), + pytest.param( + "_flash_3_varlen_hub", + marks=pytest.mark.skipif(not is_kernels_available(), reason="`kernels` is not available."), + ), ], ) @pytest.mark.parametrize("ulysses_anything", [True, False]) @@ -443,7 +447,7 @@ def test_context_parallel_attn_backend_inference(self, cp_type, attention_backen if cp_type == "ring_degree": if attention_backend == AttentionBackendName.NATIVE: pytest.skip("Skipping test because ring isn't supported with native attention backend.") - elif attention_backend in ("flash_varlen_hub"): + elif attention_backend in ("flash_varlen_hub", "_flash_3_varlen_hub"): pytest.skip("`ring_degree` is not yet supported for varlen attention hub kernels.") if ulysses_anything and "ulysses" not in cp_type: @@ -453,6 +457,18 @@ def test_context_parallel_attn_backend_inference(self, cp_type, attention_backen init_dict = self.get_init_dict() inputs_dict = self.get_dummy_inputs() + # Single-GPU reference with the same attention backend (no context parallel) + model = self.model_class(**init_dict).eval().to(torch_device) + if attention_backend: + model.set_attention_backend(attention_backend) + + # Copy inputs and cast model + inputs as needed + ref_inputs = inputs_dict.copy() + model, ref_inputs = _maybe_cast_to_bf16(attention_backend, model, ref_inputs) + state_dict = {k: v.cpu() for k, v in model.state_dict().items()} + with torch.no_grad(): + ref_output = model(**ref_inputs, return_dict=False)[0].cpu() + # Move all tensors to CPU for multiprocessing inputs_dict = {k: v.cpu() if isinstance(v, torch.Tensor) else v for k, v in inputs_dict.items()} cp_dict = {cp_type: world_size} @@ -478,6 +494,7 @@ def test_context_parallel_attn_backend_inference(self, cp_type, attention_backen inputs_dict, return_dict, attention_backend, + state_dict, ), nprocs=world_size, join=True, @@ -486,3 +503,6 @@ def test_context_parallel_attn_backend_inference(self, cp_type, attention_backen assert return_dict.get("status") == "success", ( f"Context parallel inference failed: {return_dict.get('error', 'Unknown error')}" ) + + cp_output = torch.tensor(return_dict["output"], dtype=ref_output.dtype) + torch.testing.assert_close(ref_output, cp_output, atol=1e-2, rtol=1e-2) diff --git a/tests/models/testing_utils/quantization.py b/tests/models/testing_utils/quantization.py index 4849e28fb396..a72a260d63e8 100644 --- a/tests/models/testing_utils/quantization.py +++ b/tests/models/testing_utils/quantization.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/testing_utils/single_file.py b/tests/models/testing_utils/single_file.py index 2e2d7a435fb8..6e9443aa088c 100644 --- a/tests/models/testing_utils/single_file.py +++ b/tests/models/testing_utils/single_file.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/testing_utils/training.py b/tests/models/testing_utils/training.py index 44cce6af68e5..a7a25a7cc706 100644 --- a/tests/models/testing_utils/training.py +++ b/tests/models/testing_utils/training.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -210,11 +210,16 @@ def test_mixed_precision_training(self): # Test with bfloat16 if torch.device(torch_device).type != "cpu": - model.zero_grad() - with torch.amp.autocast(device_type=torch.device(torch_device).type, dtype=torch.bfloat16): - output = model(**inputs_dict, return_dict=False)[0] - - noise = torch.randn((output.shape[0],) + self.output_shape).to(torch_device) - loss = torch.nn.functional.mse_loss(output, noise) - - loss.backward() + if torch.device(torch_device).type == "cuda" and not torch.cuda.is_bf16_supported( + including_emulation=False + ): + pytest.skip("bfloat16 training is not supported on this GPU.") + else: + model.zero_grad() + with torch.amp.autocast(device_type=torch.device(torch_device).type, dtype=torch.bfloat16): + output = model(**inputs_dict, return_dict=False)[0] + + noise = torch.randn((output.shape[0],) + self.output_shape).to(torch_device) + loss = torch.nn.functional.mse_loss(output, noise) + + loss.backward() diff --git a/tests/models/testing_utils/utils.py b/tests/models/testing_utils/utils.py index 8877755377e0..3beb59ed1a66 100644 --- a/tests/models/testing_utils/utils.py +++ b/tests/models/testing_utils/utils.py @@ -8,6 +8,7 @@ AttentionBackendName.FLASH_HUB, AttentionBackendName.FLASH_VARLEN_HUB, AttentionBackendName._FLASH_3_HUB, + AttentionBackendName._FLASH_3_VARLEN_HUB, } diff --git a/tests/models/transformers/test_models_dit_transformer2d.py b/tests/models/transformers/test_models_dit_transformer2d.py index f1efb362d104..43fbfe940e1f 100644 --- a/tests/models/transformers/test_models_dit_transformer2d.py +++ b/tests/models/transformers/test_models_dit_transformer2d.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/transformers/test_models_pixart_transformer2d.py b/tests/models/transformers/test_models_pixart_transformer2d.py index 879274c52dbd..964cbbce92c9 100644 --- a/tests/models/transformers/test_models_pixart_transformer2d.py +++ b/tests/models/transformers/test_models_pixart_transformer2d.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/transformers/test_models_prior.py b/tests/models/transformers/test_models_prior.py index 1da32b77786a..c2ebfc0ee841 100644 --- a/tests/models/transformers/test_models_prior.py +++ b/tests/models/transformers/test_models_prior.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/transformers/test_models_transformer_ace_step.py b/tests/models/transformers/test_models_transformer_ace_step.py index ba0c7769692b..d1bde6c50cd5 100644 --- a/tests/models/transformers/test_models_transformer_ace_step.py +++ b/tests/models/transformers/test_models_transformer_ace_step.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/transformers/test_models_transformer_allegro.py b/tests/models/transformers/test_models_transformer_allegro.py index 0c3e302a3f0d..9ab56af9d7c6 100644 --- a/tests/models/transformers/test_models_transformer_allegro.py +++ b/tests/models/transformers/test_models_transformer_allegro.py @@ -1,4 +1,4 @@ -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/transformers/test_models_transformer_anyflow.py b/tests/models/transformers/test_models_transformer_anyflow.py index df72567a7455..5011222f17c9 100644 --- a/tests/models/transformers/test_models_transformer_anyflow.py +++ b/tests/models/transformers/test_models_transformer_anyflow.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import pytest import torch from diffusers import AnyFlowTransformer3DModel @@ -100,12 +99,6 @@ def get_dummy_inputs(self) -> dict[str, torch.Tensor]: class TestAnyFlowTransformer3D(AnyFlowTransformer3DTesterConfig, ModelTesterMixin): """Core model tests for AnyFlow Transformer 3D (bidirectional variant).""" - @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16], ids=["fp16", "bf16"]) - def test_from_save_pretrained_dtype_inference(self, tmp_path, dtype): - # Skip: fp16/bf16 require very high atol to pass, providing little signal. - # Dtype preservation is already tested by test_from_save_pretrained_dtype and test_keep_in_fp32_modules. - pytest.skip("Tolerance requirements too high for meaningful test") - class TestAnyFlowTransformer3DMemory(AnyFlowTransformer3DTesterConfig, MemoryTesterMixin): """Memory optimization tests for AnyFlow Transformer 3D.""" diff --git a/tests/models/transformers/test_models_transformer_anyflow_far.py b/tests/models/transformers/test_models_transformer_anyflow_far.py index d7ed471fa875..b1b9d155b752 100644 --- a/tests/models/transformers/test_models_transformer_anyflow_far.py +++ b/tests/models/transformers/test_models_transformer_anyflow_far.py @@ -113,12 +113,6 @@ def get_dummy_inputs(self) -> dict[str, "torch.Tensor"]: class TestAnyFlowFARTransformer3D(AnyFlowFARTransformer3DTesterConfig, ModelTesterMixin): """Core model tests for AnyFlow FAR causal Transformer 3D.""" - @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16], ids=["fp16", "bf16"]) - def test_from_save_pretrained_dtype_inference(self, tmp_path, dtype): - # Skip: fp16/bf16 require very high atol to pass, providing little signal. - # Dtype preservation is already tested by test_from_save_pretrained_dtype and test_keep_in_fp32_modules. - pytest.skip("Tolerance requirements too high for meaningful test") - class TestAnyFlowFARTransformer3DMemory(AnyFlowFARTransformer3DTesterConfig, MemoryTesterMixin): """Memory optimization tests for AnyFlow FAR Transformer 3D.""" diff --git a/tests/models/transformers/test_models_transformer_aura_flow.py b/tests/models/transformers/test_models_transformer_aura_flow.py index 3e13945977fd..81e5920a29ba 100644 --- a/tests/models/transformers/test_models_transformer_aura_flow.py +++ b/tests/models/transformers/test_models_transformer_aura_flow.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/transformers/test_models_transformer_bria.py b/tests/models/transformers/test_models_transformer_bria.py index db4998e652b8..0f3c89e41822 100644 --- a/tests/models/transformers/test_models_transformer_bria.py +++ b/tests/models/transformers/test_models_transformer_bria.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/transformers/test_models_transformer_bria_fibo.py b/tests/models/transformers/test_models_transformer_bria_fibo.py index 6dcf725d0b28..c494c38b147b 100644 --- a/tests/models/transformers/test_models_transformer_bria_fibo.py +++ b/tests/models/transformers/test_models_transformer_bria_fibo.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/transformers/test_models_transformer_chroma.py b/tests/models/transformers/test_models_transformer_chroma.py index dac0494bf0c9..d63aee32cbf6 100644 --- a/tests/models/transformers/test_models_transformer_chroma.py +++ b/tests/models/transformers/test_models_transformer_chroma.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/transformers/test_models_transformer_chronoedit.py b/tests/models/transformers/test_models_transformer_chronoedit.py index 8baca5091b98..83878c0a9b68 100644 --- a/tests/models/transformers/test_models_transformer_chronoedit.py +++ b/tests/models/transformers/test_models_transformer_chronoedit.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/transformers/test_models_transformer_cogvideox.py b/tests/models/transformers/test_models_transformer_cogvideox.py index 97ac1b40621f..a4e80c37413f 100644 --- a/tests/models/transformers/test_models_transformer_cogvideox.py +++ b/tests/models/transformers/test_models_transformer_cogvideox.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/transformers/test_models_transformer_cogview3plus.py b/tests/models/transformers/test_models_transformer_cogview3plus.py index 97ac28a108e1..5bbcf664824e 100644 --- a/tests/models/transformers/test_models_transformer_cogview3plus.py +++ b/tests/models/transformers/test_models_transformer_cogview3plus.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/transformers/test_models_transformer_cogview4.py b/tests/models/transformers/test_models_transformer_cogview4.py index 0f390cb356e9..10497b883d4c 100644 --- a/tests/models/transformers/test_models_transformer_cogview4.py +++ b/tests/models/transformers/test_models_transformer_cogview4.py @@ -1,4 +1,4 @@ -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/transformers/test_models_transformer_consisid.py b/tests/models/transformers/test_models_transformer_consisid.py index cb02e8a359b3..efd43f6457ea 100644 --- a/tests/models/transformers/test_models_transformer_consisid.py +++ b/tests/models/transformers/test_models_transformer_consisid.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/transformers/test_models_transformer_cosmos.py b/tests/models/transformers/test_models_transformer_cosmos.py index 457d4d63f410..7698d189a5e2 100644 --- a/tests/models/transformers/test_models_transformer_cosmos.py +++ b/tests/models/transformers/test_models_transformer_cosmos.py @@ -1,4 +1,4 @@ -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/transformers/test_models_transformer_easyanimate.py b/tests/models/transformers/test_models_transformer_easyanimate.py index fe67932d1389..d7449d13c67c 100644 --- a/tests/models/transformers/test_models_transformer_easyanimate.py +++ b/tests/models/transformers/test_models_transformer_easyanimate.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/transformers/test_models_transformer_ernie_image.py b/tests/models/transformers/test_models_transformer_ernie_image.py index bff0894df08b..aff444c6941e 100644 --- a/tests/models/transformers/test_models_transformer_ernie_image.py +++ b/tests/models/transformers/test_models_transformer_ernie_image.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/transformers/test_models_transformer_flux.py b/tests/models/transformers/test_models_transformer_flux.py index be14e2d0a3ad..0214c1f65cab 100644 --- a/tests/models/transformers/test_models_transformer_flux.py +++ b/tests/models/transformers/test_models_transformer_flux.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/transformers/test_models_transformer_flux2.py b/tests/models/transformers/test_models_transformer_flux2.py index 77b5f1b86e59..9546fdb5d969 100644 --- a/tests/models/transformers/test_models_transformer_flux2.py +++ b/tests/models/transformers/test_models_transformer_flux2.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/transformers/test_models_transformer_glm_image.py b/tests/models/transformers/test_models_transformer_glm_image.py index 18510e530ab1..c1b26219683b 100644 --- a/tests/models/transformers/test_models_transformer_glm_image.py +++ b/tests/models/transformers/test_models_transformer_glm_image.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/transformers/test_models_transformer_helios.py b/tests/models/transformers/test_models_transformer_helios.py index c365c258e596..edea08127d33 100644 --- a/tests/models/transformers/test_models_transformer_helios.py +++ b/tests/models/transformers/test_models_transformer_helios.py @@ -1,4 +1,4 @@ -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -135,12 +135,6 @@ def get_dummy_inputs(self) -> dict[str, torch.Tensor]: class TestHeliosTransformer3D(HeliosTransformer3DTesterConfig, ModelTesterMixin): """Core model tests for Helios Transformer 3D.""" - @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16], ids=["fp16", "bf16"]) - def test_from_save_pretrained_dtype_inference(self, tmp_path, dtype): - # Skip: fp16/bf16 require very high atol to pass, providing little signal. - # Dtype preservation is already tested by test_from_save_pretrained_dtype and test_keep_in_fp32_modules. - pytest.skip("Tolerance requirements too high for meaningful test") - class TestHeliosTransformer3DMemory(HeliosTransformer3DTesterConfig, MemoryTesterMixin): """Memory optimization tests for Helios Transformer 3D.""" diff --git a/tests/models/transformers/test_models_transformer_hidream.py b/tests/models/transformers/test_models_transformer_hidream.py index d61811d3d1a0..c789850c0174 100644 --- a/tests/models/transformers/test_models_transformer_hidream.py +++ b/tests/models/transformers/test_models_transformer_hidream.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/transformers/test_models_transformer_hunyuan_1_5.py b/tests/models/transformers/test_models_transformer_hunyuan_1_5.py index 02eec91a1db5..7842e2e570d2 100644 --- a/tests/models/transformers/test_models_transformer_hunyuan_1_5.py +++ b/tests/models/transformers/test_models_transformer_hunyuan_1_5.py @@ -1,4 +1,4 @@ -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/transformers/test_models_transformer_hunyuan_dit.py b/tests/models/transformers/test_models_transformer_hunyuan_dit.py index 370033ef319f..eb8728976153 100644 --- a/tests/models/transformers/test_models_transformer_hunyuan_dit.py +++ b/tests/models/transformers/test_models_transformer_hunyuan_dit.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/transformers/test_models_transformer_hunyuan_video.py b/tests/models/transformers/test_models_transformer_hunyuan_video.py index cc934be125aa..145ad4b42481 100644 --- a/tests/models/transformers/test_models_transformer_hunyuan_video.py +++ b/tests/models/transformers/test_models_transformer_hunyuan_video.py @@ -1,4 +1,4 @@ -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/transformers/test_models_transformer_hunyuan_video_framepack.py b/tests/models/transformers/test_models_transformer_hunyuan_video_framepack.py index 272b7145326d..07afb75ba6dc 100644 --- a/tests/models/transformers/test_models_transformer_hunyuan_video_framepack.py +++ b/tests/models/transformers/test_models_transformer_hunyuan_video_framepack.py @@ -1,4 +1,4 @@ -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/transformers/test_models_transformer_ideogram4.py b/tests/models/transformers/test_models_transformer_ideogram4.py index d8e7318d501d..0d29f507b7e1 100644 --- a/tests/models/transformers/test_models_transformer_ideogram4.py +++ b/tests/models/transformers/test_models_transformer_ideogram4.py @@ -13,7 +13,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import pytest import torch from diffusers import Ideogram4Transformer2DModel @@ -142,14 +141,6 @@ def get_dummy_inputs(self) -> dict[str, torch.Tensor]: class TestIdeogram4Transformer(Ideogram4TransformerTesterConfig, ModelTesterMixin): """Core model tests for Ideogram 4 Transformer.""" - @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16], ids=["fp16", "bf16"]) - def test_from_save_pretrained_dtype_inference(self, tmp_path, dtype): - # Skip: the non-persistent fp32 RoPE inv_freq buffer is truncated to fp16 by the in-memory - # .to(dtype) path but kept fp32 by from_pretrained, so the two outputs diverge well beyond any - # meaningful tolerance. Dtype preservation is already covered by test_from_save_pretrained_dtype - # and test_keep_in_fp32_modules. - pytest.skip("Tolerance requirements too high for meaningful test") - class TestIdeogram4TransformerMemory(Ideogram4TransformerTesterConfig, MemoryTesterMixin): """Memory optimization tests for Ideogram 4 Transformer.""" diff --git a/tests/models/transformers/test_models_transformer_joyimage.py b/tests/models/transformers/test_models_transformer_joyimage.py index c464a44c29b5..45d15b2d470a 100644 --- a/tests/models/transformers/test_models_transformer_joyimage.py +++ b/tests/models/transformers/test_models_transformer_joyimage.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import pytest import torch from diffusers import JoyImageEditTransformer3DModel @@ -86,9 +85,7 @@ def get_dummy_inputs(self) -> dict[str, torch.Tensor]: class TestJoyImageEditTransformer(JoyImageEditTransformerTesterConfig, ModelTesterMixin): - @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16], ids=["fp16", "bf16"]) - def test_from_save_pretrained_dtype_inference(self, tmp_path, dtype): - pytest.skip("Tolerance requirements too high for meaningful test") + pass class TestJoyImageEditTransformerMemory(JoyImageEditTransformerTesterConfig, MemoryTesterMixin): diff --git a/tests/models/transformers/test_models_transformer_joyimage_edit_plus.py b/tests/models/transformers/test_models_transformer_joyimage_edit_plus.py new file mode 100644 index 000000000000..b389b6f56bbd --- /dev/null +++ b/tests/models/transformers/test_models_transformer_joyimage_edit_plus.py @@ -0,0 +1,111 @@ +# Copyright 2025 The HuggingFace Team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch + +from diffusers import JoyImageEditPlusTransformer3DModel +from diffusers.utils.torch_utils import randn_tensor + +from ...testing_utils import enable_full_determinism, torch_device +from ..testing_utils import ( + AttentionTesterMixin, + BaseModelTesterConfig, + MemoryTesterMixin, + ModelTesterMixin, + TorchCompileTesterMixin, + TrainingTesterMixin, +) + + +enable_full_determinism() + + +class JoyImageEditPlusTransformerTesterConfig(BaseModelTesterConfig): + @property + def model_class(self): + return JoyImageEditPlusTransformer3DModel + + @property + def output_shape(self) -> tuple[int, ...]: + return (2, 16, 1, 2, 2) + + @property + def input_shape(self) -> tuple[int, ...]: + return (2, 16, 1, 2, 2) + + @property + def main_input_name(self) -> str: + return "hidden_states" + + @property + def uses_custom_attn_processor(self) -> bool: + return True + + @property + def model_split_percents(self) -> list: + return [0.7, 0.6, 0.6] + + @property + def generator(self): + return torch.Generator("cpu").manual_seed(0) + + def get_init_dict(self) -> dict[str, int | list[int]]: + return { + "patch_size": [1, 2, 2], + "in_channels": 16, + "hidden_size": 32, + "num_attention_heads": 2, + "text_dim": 16, + "num_layers": 2, + "rope_dim_list": [4, 6, 6], + "theta": 256, + } + + def get_dummy_inputs(self) -> dict[str, torch.Tensor]: + batch_size = 1 + max_patches = 2 + hidden_states = randn_tensor( + (batch_size, max_patches, 16, 1, 2, 2), generator=self.generator, device=torch_device + ) + encoder_hidden_states = randn_tensor((batch_size, 12, 16), generator=self.generator, device=torch_device) + timestep = torch.tensor([1.0]).to(torch_device).expand(batch_size) + shape_list = [[(1, 1, 1), (1, 1, 1)]] + return { + "hidden_states": hidden_states, + "encoder_hidden_states": encoder_hidden_states, + "timestep": timestep, + "shape_list": shape_list, + } + + +class TestJoyImageEditPlusTransformer(JoyImageEditPlusTransformerTesterConfig, ModelTesterMixin): + pass + + +class TestJoyImageEditPlusTransformerMemory(JoyImageEditPlusTransformerTesterConfig, MemoryTesterMixin): + pass + + +class TestJoyImageEditPlusTransformerTraining(JoyImageEditPlusTransformerTesterConfig, TrainingTesterMixin): + def test_gradient_checkpointing_is_applied(self): + expected_set = {"JoyImageEditPlusTransformer3DModel"} + super().test_gradient_checkpointing_is_applied(expected_set=expected_set) + + +class TestJoyImageEditPlusTransformerAttention(JoyImageEditPlusTransformerTesterConfig, AttentionTesterMixin): + pass + + +class TestJoyImageEditPlusTransformerCompile(JoyImageEditPlusTransformerTesterConfig, TorchCompileTesterMixin): + pass diff --git a/tests/models/transformers/test_models_transformer_latte.py b/tests/models/transformers/test_models_transformer_latte.py index 946e5ce8a5a9..6a087e39ccee 100644 --- a/tests/models/transformers/test_models_transformer_latte.py +++ b/tests/models/transformers/test_models_transformer_latte.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/transformers/test_models_transformer_longcat_audio_dit.py b/tests/models/transformers/test_models_transformer_longcat_audio_dit.py index b418a3068449..c3e1575fbded 100644 --- a/tests/models/transformers/test_models_transformer_longcat_audio_dit.py +++ b/tests/models/transformers/test_models_transformer_longcat_audio_dit.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/transformers/test_models_transformer_longcat_image.py b/tests/models/transformers/test_models_transformer_longcat_image.py index 5057eda1229c..5ee446ce5742 100644 --- a/tests/models/transformers/test_models_transformer_longcat_image.py +++ b/tests/models/transformers/test_models_transformer_longcat_image.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/transformers/test_models_transformer_ltx.py b/tests/models/transformers/test_models_transformer_ltx.py index 95bebd8c3335..48397030633e 100644 --- a/tests/models/transformers/test_models_transformer_ltx.py +++ b/tests/models/transformers/test_models_transformer_ltx.py @@ -1,4 +1,4 @@ -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/transformers/test_models_transformer_ltx2.py b/tests/models/transformers/test_models_transformer_ltx2.py index e0e858bb6916..63a84e60b5d0 100644 --- a/tests/models/transformers/test_models_transformer_ltx2.py +++ b/tests/models/transformers/test_models_transformer_ltx2.py @@ -1,4 +1,4 @@ -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/transformers/test_models_transformer_lumina.py b/tests/models/transformers/test_models_transformer_lumina.py index 6c6403727288..81f2b46b9752 100644 --- a/tests/models/transformers/test_models_transformer_lumina.py +++ b/tests/models/transformers/test_models_transformer_lumina.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/transformers/test_models_transformer_lumina2.py b/tests/models/transformers/test_models_transformer_lumina2.py index 4fdd81b8762d..7e50902392f4 100644 --- a/tests/models/transformers/test_models_transformer_lumina2.py +++ b/tests/models/transformers/test_models_transformer_lumina2.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/transformers/test_models_transformer_mochi.py b/tests/models/transformers/test_models_transformer_mochi.py index 992df7cf6905..8530dd98cf90 100644 --- a/tests/models/transformers/test_models_transformer_mochi.py +++ b/tests/models/transformers/test_models_transformer_mochi.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/transformers/test_models_transformer_motif_video.py b/tests/models/transformers/test_models_transformer_motif_video.py index 8d8693acda37..ce9711f9a71d 100644 --- a/tests/models/transformers/test_models_transformer_motif_video.py +++ b/tests/models/transformers/test_models_transformer_motif_video.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/transformers/test_models_transformer_nucleusmoe_image.py b/tests/models/transformers/test_models_transformer_nucleusmoe_image.py index 14fd51d9b8c8..f4ecb9a2c62a 100644 --- a/tests/models/transformers/test_models_transformer_nucleusmoe_image.py +++ b/tests/models/transformers/test_models_transformer_nucleusmoe_image.py @@ -1,4 +1,4 @@ -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/transformers/test_models_transformer_omnigen.py b/tests/models/transformers/test_models_transformer_omnigen.py index 60d4bac84bd2..9b58056bf06d 100644 --- a/tests/models/transformers/test_models_transformer_omnigen.py +++ b/tests/models/transformers/test_models_transformer_omnigen.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/transformers/test_models_transformer_ovis_image.py b/tests/models/transformers/test_models_transformer_ovis_image.py index d7867fe56d29..a859e80b4f0f 100644 --- a/tests/models/transformers/test_models_transformer_ovis_image.py +++ b/tests/models/transformers/test_models_transformer_ovis_image.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/transformers/test_models_transformer_prx.py b/tests/models/transformers/test_models_transformer_prx.py index ad4ee07f3898..7cc6606f8379 100644 --- a/tests/models/transformers/test_models_transformer_prx.py +++ b/tests/models/transformers/test_models_transformer_prx.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/transformers/test_models_transformer_qwenimage.py b/tests/models/transformers/test_models_transformer_qwenimage.py index 33bccf816c78..365a170185f4 100644 --- a/tests/models/transformers/test_models_transformer_qwenimage.py +++ b/tests/models/transformers/test_models_transformer_qwenimage.py @@ -1,4 +1,4 @@ -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -22,6 +22,7 @@ from ...testing_utils import enable_full_determinism, torch_device from ..testing_utils import ( + AttentionBackendTesterMixin, AttentionTesterMixin, BaseModelTesterConfig, BitsAndBytesTesterMixin, @@ -250,6 +251,21 @@ class TestQwenImageTransformerAttention(QwenImageTransformerTesterConfig, Attent """Attention processor tests for QwenImage Transformer.""" +class TestQwenImageTransformerAttentionBackend(QwenImageTransformerTesterConfig, AttentionBackendTesterMixin): + """Attention backend tests for QwenImage Transformer.""" + + unsupported_attn_backends = ["flash_hub", "_flash_3_hub"] + + def get_dummy_inputs(self, batch_size: int = 2): + inputs = super().get_dummy_inputs(batch_size=batch_size) + mask = inputs["encoder_hidden_states_mask"] + mask[...] = 0 + mask[0, :2] = 1 + mask[1, :6] = 1 + inputs["encoder_hidden_states_mask"] = mask.bool() + return inputs + + class TestQwenImageTransformerContextParallel(QwenImageTransformerTesterConfig, ContextParallelTesterMixin): """Context Parallel inference tests for QwenImage Transformer.""" diff --git a/tests/models/transformers/test_models_transformer_sana.py b/tests/models/transformers/test_models_transformer_sana.py index bb585a21d56f..41d581c29090 100644 --- a/tests/models/transformers/test_models_transformer_sana.py +++ b/tests/models/transformers/test_models_transformer_sana.py @@ -1,4 +1,4 @@ -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/transformers/test_models_transformer_sana_video.py b/tests/models/transformers/test_models_transformer_sana_video.py index e9d3a2d8da8e..97f26abe1b36 100644 --- a/tests/models/transformers/test_models_transformer_sana_video.py +++ b/tests/models/transformers/test_models_transformer_sana_video.py @@ -1,4 +1,4 @@ -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/transformers/test_models_transformer_sd3.py b/tests/models/transformers/test_models_transformer_sd3.py index 50590fab8be1..e38c7853a613 100644 --- a/tests/models/transformers/test_models_transformer_sd3.py +++ b/tests/models/transformers/test_models_transformer_sd3.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/transformers/test_models_transformer_temporal.py b/tests/models/transformers/test_models_transformer_temporal.py index ff917f65cf33..48a17a7de0b4 100644 --- a/tests/models/transformers/test_models_transformer_temporal.py +++ b/tests/models/transformers/test_models_transformer_temporal.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/transformers/test_models_transformer_wan.py b/tests/models/transformers/test_models_transformer_wan.py index 60bba9dfbe18..c4c64eabcca7 100644 --- a/tests/models/transformers/test_models_transformer_wan.py +++ b/tests/models/transformers/test_models_transformer_wan.py @@ -1,4 +1,4 @@ -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import pytest import torch from diffusers import WanTransformer3DModel @@ -106,12 +105,6 @@ def get_dummy_inputs(self) -> dict[str, torch.Tensor]: class TestWanTransformer3D(WanTransformer3DTesterConfig, ModelTesterMixin): """Core model tests for Wan Transformer 3D.""" - @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16], ids=["fp16", "bf16"]) - def test_from_save_pretrained_dtype_inference(self, tmp_path, dtype): - # Skip: fp16/bf16 require very high atol to pass, providing little signal. - # Dtype preservation is already tested by test_from_save_pretrained_dtype and test_keep_in_fp32_modules. - pytest.skip("Tolerance requirements too high for meaningful test") - class TestWanTransformer3DMemory(WanTransformer3DTesterConfig, MemoryTesterMixin): """Memory optimization tests for Wan Transformer 3D.""" diff --git a/tests/models/transformers/test_models_transformer_wan_animate.py b/tests/models/transformers/test_models_transformer_wan_animate.py index bd751974637b..bcba03166a99 100644 --- a/tests/models/transformers/test_models_transformer_wan_animate.py +++ b/tests/models/transformers/test_models_transformer_wan_animate.py @@ -1,4 +1,4 @@ -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -152,12 +152,6 @@ def test_output(self, base_model_output): expected_output_shape = (1, 4, 21, 16, 16) super().test_output(base_model_output, expected_output_shape=expected_output_shape) - @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16], ids=["fp16", "bf16"]) - def test_from_save_pretrained_dtype_inference(self, tmp_path, dtype): - # Skip: fp16/bf16 require very high atol (~1e-2) to pass, providing little signal. - # Dtype preservation is already tested by test_from_save_pretrained_dtype and test_keep_in_fp32_modules. - pytest.skip("Tolerance requirements too high for meaningful test") - class TestWanAnimateTransformer3DMemory(WanAnimateTransformer3DTesterConfig, MemoryTesterMixin): """Memory optimization tests for Wan Animate Transformer 3D.""" diff --git a/tests/models/transformers/test_models_transformer_wan_vace.py b/tests/models/transformers/test_models_transformer_wan_vace.py index 1cc829f88b9d..ac980b1ea7c0 100644 --- a/tests/models/transformers/test_models_transformer_wan_vace.py +++ b/tests/models/transformers/test_models_transformer_wan_vace.py @@ -1,4 +1,4 @@ -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -117,12 +117,6 @@ def get_dummy_inputs(self) -> dict[str, torch.Tensor]: class TestWanVACETransformer3D(WanVACETransformer3DTesterConfig, ModelTesterMixin): """Core model tests for Wan VACE Transformer 3D.""" - @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16], ids=["fp16", "bf16"]) - def test_from_save_pretrained_dtype_inference(self, tmp_path, dtype): - # Skip: fp16/bf16 require very high atol to pass, providing little signal. - # Dtype preservation is already tested by test_from_save_pretrained_dtype and test_keep_in_fp32_modules. - pytest.skip("Tolerance requirements too high for meaningful test") - def test_model_parallelism(self, tmp_path): # Skip: Device mismatch between cuda:0 and cuda:1 in VACE control flow pytest.skip("Model parallelism not yet supported for WanVACE") diff --git a/tests/models/transformers/test_models_transformer_z_image.py b/tests/models/transformers/test_models_transformer_z_image.py index ad4a081557c5..85aeb34c25c4 100644 --- a/tests/models/transformers/test_models_transformer_z_image.py +++ b/tests/models/transformers/test_models_transformer_z_image.py @@ -1,4 +1,4 @@ -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/unets/test_models_unet_1d.py b/tests/models/unets/test_models_unet_1d.py index ee11cba3dd2c..913cba3e0a3b 100644 --- a/tests/models/unets/test_models_unet_1d.py +++ b/tests/models/unets/test_models_unet_1d.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/unets/test_models_unet_2d.py b/tests/models/unets/test_models_unet_2d.py index 0399f4301214..d55eff428541 100644 --- a/tests/models/unets/test_models_unet_2d.py +++ b/tests/models/unets/test_models_unet_2d.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/unets/test_models_unet_2d_condition.py b/tests/models/unets/test_models_unet_2d_condition.py index 1ac3830a6ae5..fc50a0a46125 100644 --- a/tests/models/unets/test_models_unet_2d_condition.py +++ b/tests/models/unets/test_models_unet_2d_condition.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -20,7 +20,6 @@ import unittest from collections import OrderedDict -import pytest import torch from huggingface_hub import snapshot_download from parameterized import parameterized @@ -44,9 +43,7 @@ backend_reset_peak_memory_stats, enable_full_determinism, floats_tensor, - is_peft_available, load_hf_numpy, - require_peft_backend, require_torch_accelerator, require_torch_accelerator_with_fp16, skip_mps, @@ -68,29 +65,11 @@ ) -if is_peft_available(): - from peft import LoraConfig - - from ..testing_utils.lora import check_if_lora_correctly_set - - logger = logging.get_logger(__name__) enable_full_determinism() -def get_unet_lora_config(): - rank = 4 - unet_lora_config = LoraConfig( - r=rank, - lora_alpha=rank, - target_modules=["to_q", "to_k", "to_v", "to_out.0"], - init_lora_weights=False, - use_dora=False, - ) - return unet_lora_config - - def create_ip_adapter_state_dict(model): # "ip_adapter" (cross-attention weights) ip_cross_attn_state_dict = {} @@ -694,63 +673,6 @@ def test_load_sharded_checkpoint_device_map_from_hub_local_subfolder(self): class TestUNet2DConditionLoRA(UNet2DConditionTesterConfig, LoraTesterMixin): """LoRA adapter tests for UNet2DConditionModel.""" - @require_peft_backend - def test_load_attn_procs_raise_warning(self): - """Test that deprecated load_attn_procs method raises FutureWarning.""" - init_dict = self.get_init_dict() - inputs_dict = self.get_dummy_inputs() - model = self.model_class(**init_dict) - model.to(torch_device) - - # forward pass without LoRA - with torch.no_grad(): - non_lora_sample = model(**inputs_dict).sample - - unet_lora_config = get_unet_lora_config() - model.add_adapter(unet_lora_config) - - assert check_if_lora_correctly_set(model), "Lora not correctly set in UNet." - - # forward pass with LoRA - with torch.no_grad(): - lora_sample_1 = model(**inputs_dict).sample - - with tempfile.TemporaryDirectory() as tmpdirname: - model.save_attn_procs(tmpdirname) - model.unload_lora() - - with pytest.warns(FutureWarning, match="Using the `load_attn_procs\\(\\)` method has been deprecated"): - model.load_attn_procs(os.path.join(tmpdirname, "pytorch_lora_weights.safetensors")) - - # import to still check for the rest of the stuff. - assert check_if_lora_correctly_set(model), "Lora not correctly set in UNet." - - with torch.no_grad(): - lora_sample_2 = model(**inputs_dict).sample - - assert not torch.allclose(non_lora_sample, lora_sample_1, atol=1e-4, rtol=1e-4), ( - "LoRA injected UNet should produce different results." - ) - assert torch.allclose(lora_sample_1, lora_sample_2, atol=1e-4, rtol=1e-4), ( - "Loading from a saved checkpoint should produce identical results." - ) - - @require_peft_backend - def test_save_attn_procs_raise_warning(self): - """Test that deprecated save_attn_procs method raises FutureWarning.""" - init_dict = self.get_init_dict() - model = self.model_class(**init_dict) - model.to(torch_device) - - unet_lora_config = get_unet_lora_config() - model.add_adapter(unet_lora_config) - - assert check_if_lora_correctly_set(model), "Lora not correctly set in UNet." - - with tempfile.TemporaryDirectory() as tmpdirname: - with pytest.warns(FutureWarning, match="Using the `save_attn_procs\\(\\)` method has been deprecated"): - model.save_attn_procs(os.path.join(tmpdirname)) - class TestUNet2DConditionMemory(UNet2DConditionTesterConfig, MemoryTesterMixin): """Memory optimization tests for UNet2DConditionModel.""" diff --git a/tests/models/unets/test_models_unet_3d_condition.py b/tests/models/unets/test_models_unet_3d_condition.py index 9ebdd659ee4b..bb05c8d87058 100644 --- a/tests/models/unets/test_models_unet_3d_condition.py +++ b/tests/models/unets/test_models_unet_3d_condition.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/unets/test_models_unet_controlnetxs.py b/tests/models/unets/test_models_unet_controlnetxs.py index a12eb57228da..db8433cc4652 100644 --- a/tests/models/unets/test_models_unet_controlnetxs.py +++ b/tests/models/unets/test_models_unet_controlnetxs.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/unets/test_models_unet_motion.py b/tests/models/unets/test_models_unet_motion.py index 21bc9bd62de9..e38ef7ccdad8 100644 --- a/tests/models/unets/test_models_unet_motion.py +++ b/tests/models/unets/test_models_unet_motion.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/unets/test_models_unet_spatiotemporal.py b/tests/models/unets/test_models_unet_spatiotemporal.py index fc2ad7dec8dc..24f41b67ab85 100644 --- a/tests/models/unets/test_models_unet_spatiotemporal.py +++ b/tests/models/unets/test_models_unet_spatiotemporal.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/unets/test_unet_2d_blocks.py b/tests/models/unets/test_unet_2d_blocks.py index 5c006963e30c..6e550b05a8b9 100644 --- a/tests/models/unets/test_unet_2d_blocks.py +++ b/tests/models/unets/test_unet_2d_blocks.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/models/unets/test_unet_blocks_common.py b/tests/models/unets/test_unet_blocks_common.py index 85f9bf8353bf..1e38e6ce9b23 100644 --- a/tests/models/unets/test_unet_blocks_common.py +++ b/tests/models/unets/test_unet_blocks_common.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/modular_pipelines/cosmos/__init__.py b/tests/modular_pipelines/cosmos/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py new file mode 100644 index 000000000000..5297a441cebf --- /dev/null +++ b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py @@ -0,0 +1,346 @@ +# coding=utf-8 +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json + +import pytest +import torch +from PIL import Image + +from diffusers import Cosmos3OmniTransformer, ModularPipeline +from diffusers.modular_pipelines import Cosmos3OmniBlocks, Cosmos3OmniModularPipeline +from diffusers.modular_pipelines.cosmos.before_denoise import ( + Cosmos3ActionDenoiseInputStep, + Cosmos3ActionPackSequenceStep, + Cosmos3SoundDenoiseInputStep, + Cosmos3VisionDenoiseInputStep, + Cosmos3VisionPackSequenceStep, +) +from diffusers.modular_pipelines.modular_pipeline import PipelineState + +from ...testing_utils import torch_device +from ..test_modular_pipelines_common import ModularPipelineTesterMixin + + +TEXT_VISION_WORKFLOW = [ + ("text_encoder", "Cosmos3TextEncoderStep"), + ("denoise.prepare_text_segments", "Cosmos3PrepareTextSegmentsStep"), + ("denoise.prepare_vision_latents", "Cosmos3VisionPrepareLatentsStep"), + ("denoise.pack_vision_sequence", "Cosmos3VisionPackSequenceStep"), + ("denoise.prepare_vision_denoiser_inputs", "Cosmos3VisionDenoiseInputStep"), + ("denoise.set_timesteps", "Cosmos3SetTimestepsStep"), + ("denoise.denoise", "Cosmos3VisionDenoiseStep"), + ("decode.video", "Cosmos3VideoDecodeStep"), + ("after_decode", "Cosmos3ActionOutputStep"), +] + +IMAGE_VISION_WORKFLOW = [ + ("text_encoder", "Cosmos3TextEncoderStep"), + ("vae_encoder", "Cosmos3ImageVaeEncoderStep"), + *TEXT_VISION_WORKFLOW[1:], +] + +VIDEO_VISION_WORKFLOW = [ + ("text_encoder", "Cosmos3TextEncoderStep"), + ("vae_encoder", "Cosmos3VideoVaeEncoderStep"), + *TEXT_VISION_WORKFLOW[1:], +] + +TEXT_VISION_SOUND_WORKFLOW = [ + *TEXT_VISION_WORKFLOW[:6], + ("denoise.prepare_sound_latents", "Cosmos3SoundPrepareLatentsStep"), + ("denoise.pack_sound_sequence", "Cosmos3SoundPackSequenceStep"), + ("denoise.prepare_sound_denoiser_inputs", "Cosmos3SoundDenoiseInputStep"), + ("denoise.denoise", "Cosmos3VisionSoundDenoiseStep"), + ("decode.video", "Cosmos3VideoDecodeStep"), + ("decode.sound", "Cosmos3SoundDecodeStep"), + ("after_decode", "Cosmos3ActionOutputStep"), +] + +IMAGE_VISION_SOUND_WORKFLOW = [ + ("text_encoder", "Cosmos3TextEncoderStep"), + ("vae_encoder", "Cosmos3ImageVaeEncoderStep"), + *TEXT_VISION_SOUND_WORKFLOW[1:], +] + +VIDEO_VISION_SOUND_WORKFLOW = [ + ("text_encoder", "Cosmos3TextEncoderStep"), + ("vae_encoder", "Cosmos3VideoVaeEncoderStep"), + *TEXT_VISION_SOUND_WORKFLOW[1:], +] + +ACTION_WORKFLOW = [ + ("text_encoder", "Cosmos3ActionTextStep"), + ("vae_encoder", "Cosmos3ActionVisionVaeEncoderStep"), + ("denoise.prepare_text_segments", "Cosmos3PrepareTextSegmentsStep"), + ("denoise.prepare_vision_latents", "Cosmos3VisionPrepareLatentsStep"), + ("denoise.pack_vision_sequence", "Cosmos3VisionPackSequenceStep"), + ("denoise.prepare_vision_denoiser_inputs", "Cosmos3VisionDenoiseInputStep"), + ("denoise.set_timesteps", "Cosmos3SetTimestepsStep"), + ("denoise.prepare_action_latents", "Cosmos3ActionPrepareLatentsStep"), + ("denoise.pack_action_sequence", "Cosmos3ActionPackSequenceStep"), + ("denoise.prepare_action_denoiser_inputs", "Cosmos3ActionDenoiseInputStep"), + ("denoise.denoise", "Cosmos3VisionActionDenoiseStep"), + ("decode.video", "Cosmos3VideoDecodeStep"), + ("after_decode", "Cosmos3ActionOutputStep"), +] + +COSMOS3_OMNI_WORKFLOWS = { + "text2image": TEXT_VISION_WORKFLOW, + "text2video": TEXT_VISION_WORKFLOW, + "image2video": IMAGE_VISION_WORKFLOW, + "video2video": VIDEO_VISION_WORKFLOW, + "text2video_with_sound": TEXT_VISION_SOUND_WORKFLOW, + "image2video_with_sound": IMAGE_VISION_SOUND_WORKFLOW, + "video2video_with_sound": VIDEO_VISION_SOUND_WORKFLOW, + "action_policy": ACTION_WORKFLOW, + "action_forward_dynamics": ACTION_WORKFLOW, + "action_inverse_dynamics": ACTION_WORKFLOW, +} + + +class TestCosmos3OmniModularPipelineFast(ModularPipelineTesterMixin): + pipeline_class = Cosmos3OmniModularPipeline + pipeline_blocks_class = Cosmos3OmniBlocks + pretrained_model_name_or_path = "atharvajoshi10/tiny-cosmos3-modular-pipe" + + params = frozenset(["prompt", "height", "width", "num_frames", "guidance_scale"]) + batch_params = frozenset() + optional_params = frozenset(["num_inference_steps", "output_type"]) + output_name = "videos" + expected_workflow_blocks = COSMOS3_OMNI_WORKFLOWS + + def get_pipeline(self, components_manager=None, torch_dtype=torch.float32): + pipe = super().get_pipeline(components_manager, torch_dtype) + pipe.disable_safety_checker() + return pipe + + def get_dummy_inputs(self, seed=0): + return { + "prompt": "A small robot moves across a table.", + "negative_prompt": "", + "generator": self.get_generator(seed), + "num_inference_steps": 2, + "guidance_scale": 2.0, + "height": 32, + "width": 32, + "num_frames": 5, + "output_type": "latent", + } + + @pytest.mark.skip(reason="Cosmos3 does not support batched prompts.") + def test_inference_batch_consistent(self): + pass + + @pytest.mark.skip(reason="Cosmos3 does not support batched prompts.") + def test_inference_batch_single_identical(self): + pass + + @pytest.mark.skip(reason="Cosmos3 does not support multiple videos per prompt.") + def test_num_images_per_prompt(self): + pass + + @pytest.mark.skip(reason="Cosmos3 checkpoints support bfloat16, not float16, inference.") + def test_float16_inference(self): + pass + + def test_save_from_pretrained(self, tmp_path): + base_pipe = self.get_pipeline().to(torch_device) + base_pipe.save_pretrained(str(tmp_path)) + + loaded_pipe = ModularPipeline.from_pretrained(str(tmp_path)) + loaded_pipe.load_components(torch_dtype=torch.float32) + loaded_pipe.disable_safety_checker() + loaded_pipe.to(torch_device) + + base_output = base_pipe(**self.get_dummy_inputs(), output=self.output_name) + loaded_output = loaded_pipe(**self.get_dummy_inputs(), output=self.output_name) + + assert torch.abs(base_output - loaded_output).max() < 1e-3 + + def test_vae_encoder_is_standalone_and_validates_conditioning_inputs(self): + pipe = self.get_pipeline() + vae_encoder = pipe.blocks.sub_blocks["vae_encoder"] + + assert vae_encoder.select_block(action=None, image=None, video=None) is None + assert vae_encoder.select_block(action=None, image=object(), video=None) == "image_conditioning" + assert vae_encoder.select_block(action=None, image=None, video=object()) == "video_conditioning" + assert vae_encoder.select_block(action=object(), image=None, video=None) == "action_conditioning" + + state = PipelineState() + state.set("image", Image.new("RGB", (32, 32))) + state.set("num_frames", 5) + state.set("height", 32) + state.set("width", 32) + _, state = vae_encoder(pipe, state) + + assert state.get("x0_tokens_vision") is not None + assert state.get("vision_condition_frames") == [0] + + action_vae_encoder = vae_encoder.sub_blocks["action_conditioning"] + assert [input_param.name for input_param in action_vae_encoder.inputs] == ["action"] + assert [output_param.name for output_param in action_vae_encoder.intermediate_outputs] == [ + "x0_tokens_vision", + "vision_condition_frames", + "action_condition_frame_indexes", + ] + + with pytest.raises(ValueError, match="not top-level image/video"): + pipe.blocks.get_execution_blocks(action=object(), image=object()) + with pytest.raises(ValueError, match="not top-level image/video"): + pipe.blocks.get_execution_blocks(action=object(), video=object()) + with pytest.raises(ValueError, match="either image or video"): + pipe.blocks.get_execution_blocks(image=object(), video=object()) + + inputs = self.get_dummy_inputs() + inputs.update(image=Image.new("RGB", (32, 32)), num_frames=1) + with pytest.raises(ValueError, match="image-to-image generation is not supported"): + pipe(**inputs, output=self.output_name) + + @pytest.mark.parametrize("prompt_name", ["prompt", "negative_prompt"]) + def test_rejects_batched_prompts(self, prompt_name): + pipe = self.get_pipeline() + inputs = self.get_dummy_inputs() + inputs[prompt_name] = ["first prompt", "second prompt"] + + with pytest.raises(ValueError, match="batched prompts are not supported"): + pipe(**inputs, output=self.output_name) + + +def test_cosmos3_pack_steps_do_not_require_vae(): + transformer = Cosmos3OmniTransformer( + hidden_size=32, + intermediate_size=64, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=8, + latent_channel=4, + latent_patch_size=2, + patch_latent_dim=16, + vocab_size=16, + rope_scaling={"mrope_section": [2, 1, 1]}, + action_gen=True, + action_dim=10, + ) + pipe = Cosmos3OmniModularPipeline(blocks=Cosmos3VisionPackSequenceStep()) + pipe.update_components(transformer=transformer) + + assert [component.name for component in Cosmos3VisionPackSequenceStep().expected_components] == ["transformer"] + assert [component.name for component in Cosmos3ActionPackSequenceStep().expected_components] == ["transformer"] + assert not hasattr(pipe, "vae") + + vision_segment = pipe._prepare_vision_segment( + input_vision_tokens=torch.zeros(1, transformer.config.latent_channel, 2, 2, 2), + has_image_condition=False, + mrope_offset=0, + vision_fps=5.0, + curr=1, + device="cpu", + ) + action_segment = pipe._prepare_action_segment( + input_action_tokens=torch.zeros(4, transformer.config.action_dim), + condition_frame_indexes=[], + mrope_offset=0, + action_fps=5.0, + curr=1, + device="cpu", + ) + + assert vision_segment["vision_mrope_ids"].shape == (3, 2) + assert action_segment["action_mrope_ids"].shape == (3, 4) + + +def test_cosmos3_denoise_input_steps_assemble_modality_segments(): + cond_text_segment = {"text_mrope_ids": torch.tensor([[1, 2], [3, 4], [5, 6]]), "und_len": 2} + uncond_text_segment = {"text_mrope_ids": torch.tensor([[7], [8], [9]]), "und_len": 1} + cond_vision_segment = {"vision_mrope_ids": torch.tensor([[10], [11], [12]]), "num_vision_tokens": 1} + uncond_vision_segment = {"vision_mrope_ids": torch.tensor([[13, 14], [15, 16], [17, 18]]), "num_vision_tokens": 2} + cond_sound_segment = {"sound_mrope_ids": torch.tensor([[19, 20], [21, 22], [23, 24]]), "sound_len": 2} + uncond_sound_segment = {"sound_mrope_ids": torch.tensor([[25], [26], [27]]), "sound_len": 1} + cond_action_segment = {"action_mrope_ids": torch.tensor([[28], [29], [30]]), "action_len": 1} + uncond_action_segment = {"action_mrope_ids": torch.tensor([[31, 32], [33, 34], [35, 36]]), "action_len": 2} + + state = PipelineState() + state.set("cond_text_segment", cond_text_segment) + state.set("uncond_text_segment", uncond_text_segment) + state.set("cond_vision_segment", cond_vision_segment) + state.set("uncond_vision_segment", uncond_vision_segment) + state.set("cond_sound_segment", cond_sound_segment) + state.set("uncond_sound_segment", uncond_sound_segment) + state.set("cond_action_segment", cond_action_segment) + state.set("uncond_action_segment", uncond_action_segment) + + _, state = Cosmos3VisionDenoiseInputStep()(None, state) + _, state = Cosmos3SoundDenoiseInputStep()(None, state) + _, state = Cosmos3ActionDenoiseInputStep()(None, state) + + torch.testing.assert_close( + state.get("cond_position_ids"), + torch.cat( + [ + cond_text_segment["text_mrope_ids"], + cond_vision_segment["vision_mrope_ids"], + cond_sound_segment["sound_mrope_ids"], + cond_action_segment["action_mrope_ids"], + ], + dim=1, + ), + ) + torch.testing.assert_close( + state.get("uncond_position_ids"), + torch.cat( + [ + uncond_text_segment["text_mrope_ids"], + uncond_vision_segment["vision_mrope_ids"], + uncond_sound_segment["sound_mrope_ids"], + uncond_action_segment["action_mrope_ids"], + ], + dim=1, + ), + ) + assert state.get("cond_sequence_length") == 6 + assert state.get("uncond_sequence_length") == 6 + assert state.get("cond_packed_static") is None + assert state.get("uncond_packed_static") is None + + +def test_cosmos3_modular_model_index_takes_precedence(tmp_path): + (tmp_path / "model_index.json").write_text(json.dumps({"_class_name": "Cosmos3OmniDiffusersPipeline"})) + (tmp_path / "modular_model_index.json").write_text( + json.dumps( + { + "_class_name": "Cosmos3OmniModularPipeline", + "_blocks_class_name": "Cosmos3OmniBlocks", + } + ) + ) + + modular_pipe = ModularPipeline.from_pretrained(str(tmp_path)) + explicit_modular_pipe = Cosmos3OmniModularPipeline.from_pretrained(str(tmp_path)) + + assert isinstance(modular_pipe, Cosmos3OmniModularPipeline) + assert isinstance(modular_pipe.blocks, Cosmos3OmniBlocks) + assert isinstance(explicit_modular_pipe, Cosmos3OmniModularPipeline) + assert isinstance(explicit_modular_pipe.blocks, Cosmos3OmniBlocks) + + +def test_cosmos3_model_index_fallback_resolves_modular_pipeline(tmp_path): + (tmp_path / "model_index.json").write_text(json.dumps({"_class_name": "Cosmos3OmniPipeline"})) + + pipe = ModularPipeline.from_pretrained(str(tmp_path)) + + assert isinstance(pipe, Cosmos3OmniModularPipeline) + assert isinstance(pipe.blocks, Cosmos3OmniBlocks) diff --git a/tests/modular_pipelines/ernie_image/test_modular_pipeline_ernie_image.py b/tests/modular_pipelines/ernie_image/test_modular_pipeline_ernie_image.py index 511a5dc1b3eb..0c1b82700044 100644 --- a/tests/modular_pipelines/ernie_image/test_modular_pipeline_ernie_image.py +++ b/tests/modular_pipelines/ernie_image/test_modular_pipeline_ernie_image.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/modular_pipelines/flux/test_modular_pipeline_flux.py b/tests/modular_pipelines/flux/test_modular_pipeline_flux.py index 05fe16e372ec..d82a9e2f395b 100644 --- a/tests/modular_pipelines/flux/test_modular_pipeline_flux.py +++ b/tests/modular_pipelines/flux/test_modular_pipeline_flux.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/modular_pipelines/flux2/test_modular_pipeline_flux2.py b/tests/modular_pipelines/flux2/test_modular_pipeline_flux2.py index 3045af636841..2acb2d4f6dae 100644 --- a/tests/modular_pipelines/flux2/test_modular_pipeline_flux2.py +++ b/tests/modular_pipelines/flux2/test_modular_pipeline_flux2.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/modular_pipelines/flux2/test_modular_pipeline_flux2_klein.py b/tests/modular_pipelines/flux2/test_modular_pipeline_flux2_klein.py index ad295a961357..6a6984428d5d 100644 --- a/tests/modular_pipelines/flux2/test_modular_pipeline_flux2_klein.py +++ b/tests/modular_pipelines/flux2/test_modular_pipeline_flux2_klein.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/modular_pipelines/flux2/test_modular_pipeline_flux2_klein_base.py b/tests/modular_pipelines/flux2/test_modular_pipeline_flux2_klein_base.py index b3aa79040317..20328a8e4d99 100644 --- a/tests/modular_pipelines/flux2/test_modular_pipeline_flux2_klein_base.py +++ b/tests/modular_pipelines/flux2/test_modular_pipeline_flux2_klein_base.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/modular_pipelines/helios/test_modular_pipeline_helios.py b/tests/modular_pipelines/helios/test_modular_pipeline_helios.py index 44a01dad6525..cdd7cfa19fc9 100644 --- a/tests/modular_pipelines/helios/test_modular_pipeline_helios.py +++ b/tests/modular_pipelines/helios/test_modular_pipeline_helios.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/modular_pipelines/hunyuan_video1_5/test_modular_pipeline_hunyuan_video1_5.py b/tests/modular_pipelines/hunyuan_video1_5/test_modular_pipeline_hunyuan_video1_5.py index 6d776eaa1a11..a6db1289bb35 100644 --- a/tests/modular_pipelines/hunyuan_video1_5/test_modular_pipeline_hunyuan_video1_5.py +++ b/tests/modular_pipelines/hunyuan_video1_5/test_modular_pipeline_hunyuan_video1_5.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/modular_pipelines/ltx/test_modular_pipeline_ltx.py b/tests/modular_pipelines/ltx/test_modular_pipeline_ltx.py index 97efd4dd9698..733b79e98de6 100644 --- a/tests/modular_pipelines/ltx/test_modular_pipeline_ltx.py +++ b/tests/modular_pipelines/ltx/test_modular_pipeline_ltx.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/modular_pipelines/qwen/test_modular_pipeline_qwenimage.py b/tests/modular_pipelines/qwen/test_modular_pipeline_qwenimage.py index 92573c202e49..71345ace1d6b 100644 --- a/tests/modular_pipelines/qwen/test_modular_pipeline_qwenimage.py +++ b/tests/modular_pipelines/qwen/test_modular_pipeline_qwenimage.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/modular_pipelines/stable_diffusion_xl/test_modular_pipeline_stable_diffusion_xl.py b/tests/modular_pipelines/stable_diffusion_xl/test_modular_pipeline_stable_diffusion_xl.py index f640f0ec83f2..7bbfac596365 100644 --- a/tests/modular_pipelines/stable_diffusion_xl/test_modular_pipeline_stable_diffusion_xl.py +++ b/tests/modular_pipelines/stable_diffusion_xl/test_modular_pipeline_stable_diffusion_xl.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/modular_pipelines/wan/test_modular_pipeline_wan.py b/tests/modular_pipelines/wan/test_modular_pipeline_wan.py index c5ed9613e40f..1141633f2b7d 100644 --- a/tests/modular_pipelines/wan/test_modular_pipeline_wan.py +++ b/tests/modular_pipelines/wan/test_modular_pipeline_wan.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/modular_pipelines/z_image/test_modular_pipeline_z_image.py b/tests/modular_pipelines/z_image/test_modular_pipeline_z_image.py index ab45def3ef30..49de376024d6 100644 --- a/tests/modular_pipelines/z_image/test_modular_pipeline_z_image.py +++ b/tests/modular_pipelines/z_image/test_modular_pipeline_z_image.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/others/test_config.py b/tests/others/test_config.py index 232bf9d473b8..376633601f7e 100644 --- a/tests/others/test_config.py +++ b/tests/others/test_config.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/others/test_ema.py b/tests/others/test_ema.py index 436bbe1d53ff..87820ed6af84 100644 --- a/tests/others/test_ema.py +++ b/tests/others/test_ema.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/others/test_flashpack.py b/tests/others/test_flashpack.py index 85b7d09fe8db..c14410c0d8e0 100644 --- a/tests/others/test_flashpack.py +++ b/tests/others/test_flashpack.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/others/test_hub_utils.py b/tests/others/test_hub_utils.py index 0a6b8ef2bd9f..de6ab198e501 100644 --- a/tests/others/test_hub_utils.py +++ b/tests/others/test_hub_utils.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/others/test_image_processor.py b/tests/others/test_image_processor.py index e9e5c0670676..88e82ab54b82 100644 --- a/tests/others/test_image_processor.py +++ b/tests/others/test_image_processor.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/others/test_training.py b/tests/others/test_training.py index d8e86984ef1e..a339ee8a3c6b 100644 --- a/tests/others/test_training.py +++ b/tests/others/test_training.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/others/test_utils.py b/tests/others/test_utils.py index 5db007b7ed6d..412c7478c5a7 100755 --- a/tests/others/test_utils.py +++ b/tests/others/test_utils.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/others/test_video_processor.py b/tests/others/test_video_processor.py index 35c9f99c37ae..7fd98f26652f 100644 --- a/tests/others/test_video_processor.py +++ b/tests/others/test_video_processor.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/pipelines/ace_step/test_ace_step.py b/tests/pipelines/ace_step/test_ace_step.py index 6be8bfd155f0..089a9b849523 100644 --- a/tests/pipelines/ace_step/test_ace_step.py +++ b/tests/pipelines/ace_step/test_ace_step.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/pipelines/animatediff/test_animatediff.py b/tests/pipelines/animatediff/test_animatediff.py index 8d4cd4cf2c1a..0d5c2b32b9be 100644 --- a/tests/pipelines/animatediff/test_animatediff.py +++ b/tests/pipelines/animatediff/test_animatediff.py @@ -189,33 +189,33 @@ def test_ip_adapter(self): if torch_device == "cpu": expected_pipe_slice = np.array( [ - 0.5216, - 0.5620, - 0.4927, - 0.5082, - 0.4786, - 0.5932, - 0.5125, - 0.4514, - 0.5315, - 0.4694, - 0.3276, - 0.4863, - 0.3920, - 0.3684, - 0.5745, - 0.4499, - 0.5081, - 0.5414, - 0.6014, - 0.5062, - 0.3630, - 0.5296, - 0.6018, - 0.5098, - 0.4948, - 0.5101, - 0.5620, + 0.5137, + 0.5648, + 0.4867, + 0.4968, + 0.4716, + 0.5940, + 0.5136, + 0.4370, + 0.5325, + 0.4856, + 0.3607, + 0.4660, + 0.4036, + 0.3620, + 0.5743, + 0.4617, + 0.4962, + 0.5454, + 0.5908, + 0.5164, + 0.3581, + 0.5272, + 0.6035, + 0.5103, + 0.4988, + 0.5016, + 0.5651, ] ) return super().test_ip_adapter(expected_pipe_slice=expected_pipe_slice) @@ -223,7 +223,7 @@ def test_ip_adapter(self): def test_dict_tuple_outputs_equivalent(self): expected_slice = None if torch_device == "cpu": - expected_slice = np.array([0.5125, 0.4514, 0.5315, 0.4499, 0.5081, 0.5414, 0.4948, 0.5101, 0.5620]) + expected_slice = np.array([0.5136, 0.4370, 0.5325, 0.4617, 0.4962, 0.5454, 0.4988, 0.5016, 0.5651]) return super().test_dict_tuple_outputs_equivalent(expected_slice=expected_slice) def test_inference_batch_single_identical( @@ -583,7 +583,7 @@ def test_animatediff(self): steps_offset=1, clip_sample=False, ) - pipe.enable_vae_slicing() + pipe.vae.enable_slicing() pipe.enable_model_cpu_offload(device=torch_device) pipe.set_progress_bar_config(disable=None) diff --git a/tests/pipelines/animatediff/test_animatediff_controlnet.py b/tests/pipelines/animatediff/test_animatediff_controlnet.py index 4b0eb01d067c..ab354e65739d 100644 --- a/tests/pipelines/animatediff/test_animatediff_controlnet.py +++ b/tests/pipelines/animatediff/test_animatediff_controlnet.py @@ -198,24 +198,24 @@ def test_ip_adapter(self): if torch_device == "cpu": expected_pipe_slice = np.array( [ - 0.6604, - 0.4099, - 0.4928, - 0.5706, - 0.5096, - 0.5012, - 0.6051, - 0.5169, - 0.5021, - 0.4864, - 0.4261, - 0.5779, - 0.5822, - 0.4049, - 0.5253, - 0.6160, - 0.4150, - 0.5155, + 0.6680, + 0.5061, + 0.5069, + 0.5930, + 0.5747, + 0.4737, + 0.5885, + 0.5630, + 0.5083, + 0.4910, + 0.4132, + 0.5721, + 0.5793, + 0.4540, + 0.5094, + 0.5943, + 0.4598, + 0.5104, ] ) return super().test_ip_adapter(expected_pipe_slice=expected_pipe_slice) @@ -223,7 +223,7 @@ def test_ip_adapter(self): def test_dict_tuple_outputs_equivalent(self): expected_slice = None if torch_device == "cpu": - expected_slice = np.array([0.6051, 0.5169, 0.5021, 0.6160, 0.4150, 0.5155]) + expected_slice = np.array([0.5885, 0.5630, 0.5083, 0.5943, 0.4598, 0.5104]) return super().test_dict_tuple_outputs_equivalent(expected_slice=expected_slice) def test_inference_batch_single_identical( @@ -510,7 +510,7 @@ def test_vae_slicing(self, video_count=2): output_1 = pipe(**inputs) # make sure sliced vae decode yields the same result - pipe.enable_vae_slicing() + pipe.vae.enable_slicing() inputs = self.get_dummy_inputs(device) inputs["prompt"] = [inputs["prompt"]] * video_count inputs["conditioning_frames"] = [inputs["conditioning_frames"]] * video_count diff --git a/tests/pipelines/animatediff/test_animatediff_sparsectrl.py b/tests/pipelines/animatediff/test_animatediff_sparsectrl.py index 6b9f672cc4a1..c414f64d3a8b 100644 --- a/tests/pipelines/animatediff/test_animatediff_sparsectrl.py +++ b/tests/pipelines/animatediff/test_animatediff_sparsectrl.py @@ -200,24 +200,24 @@ def test_ip_adapter(self): if torch_device == "cpu": expected_pipe_slice = np.array( [ - 0.6604, - 0.4099, - 0.4928, - 0.5706, - 0.5096, - 0.5012, - 0.6051, - 0.5169, - 0.5021, - 0.4864, - 0.4261, - 0.5779, - 0.5822, - 0.4049, - 0.5253, - 0.6160, - 0.4150, - 0.5155, + 0.6680, + 0.5061, + 0.5069, + 0.5930, + 0.5747, + 0.4737, + 0.5885, + 0.5630, + 0.5083, + 0.4910, + 0.4132, + 0.5721, + 0.5793, + 0.4540, + 0.5094, + 0.5943, + 0.4598, + 0.5104, ] ) return super().test_ip_adapter(expected_pipe_slice=expected_pipe_slice) @@ -225,7 +225,7 @@ def test_ip_adapter(self): def test_dict_tuple_outputs_equivalent(self): expected_slice = None if torch_device == "cpu": - expected_slice = np.array([0.6051, 0.5169, 0.5021, 0.6160, 0.4150, 0.5155]) + expected_slice = np.array([0.5885, 0.5630, 0.5083, 0.5943, 0.4598, 0.5104]) return super().test_dict_tuple_outputs_equivalent(expected_slice=expected_slice) def test_inference_batch_single_identical( diff --git a/tests/pipelines/animatediff/test_animatediff_video2video.py b/tests/pipelines/animatediff/test_animatediff_video2video.py index 1adb13dc4cc5..deddfeff3d1e 100644 --- a/tests/pipelines/animatediff/test_animatediff_video2video.py +++ b/tests/pipelines/animatediff/test_animatediff_video2video.py @@ -182,23 +182,23 @@ def test_ip_adapter(self): expected_pipe_slice = np.array( [ 0.5569, - 0.6250, - 0.4145, - 0.5613, - 0.5563, - 0.5213, - 0.5092, - 0.4950, - 0.4950, - 0.5685, - 0.3858, - 0.4864, - 0.6458, - 0.4312, - 0.5518, - 0.5608, - 0.4418, - 0.5378, + 0.6257, + 0.4152, + 0.5620, + 0.5558, + 0.5206, + 0.5101, + 0.4937, + 0.4943, + 0.5694, + 0.3849, + 0.4863, + 0.6459, + 0.4288, + 0.5531, + 0.5620, + 0.4404, + 0.5379, ] ) return super().test_ip_adapter(expected_pipe_slice=expected_pipe_slice) diff --git a/tests/pipelines/animatediff/test_animatediff_video2video_controlnet.py b/tests/pipelines/animatediff/test_animatediff_video2video_controlnet.py index c71c8c8817dc..c69d5daaf93b 100644 --- a/tests/pipelines/animatediff/test_animatediff_video2video_controlnet.py +++ b/tests/pipelines/animatediff/test_animatediff_video2video_controlnet.py @@ -198,23 +198,23 @@ def test_ip_adapter(self): expected_pipe_slice = np.array( [ 0.5569, - 0.6250, - 0.4144, - 0.5613, - 0.5563, - 0.5213, - 0.5091, - 0.4950, - 0.4950, - 0.5684, - 0.3858, + 0.6257, + 0.4152, + 0.5620, + 0.5558, + 0.5206, + 0.5101, + 0.4937, + 0.4943, + 0.5694, + 0.3849, 0.4863, - 0.6457, - 0.4311, - 0.5517, - 0.5608, - 0.4417, - 0.5377, + 0.6459, + 0.4288, + 0.5531, + 0.5620, + 0.4404, + 0.5379, ] ) return super().test_ip_adapter(expected_pipe_slice=expected_pipe_slice) diff --git a/tests/pipelines/audioldm2/test_audioldm2.py b/tests/pipelines/audioldm2/test_audioldm2.py index 33f41bce8f16..c8aa9adeee8f 100644 --- a/tests/pipelines/audioldm2/test_audioldm2.py +++ b/tests/pipelines/audioldm2/test_audioldm2.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/pipelines/chronoedit/test_chronoedit.py b/tests/pipelines/chronoedit/test_chronoedit.py index 0b72f93eed3c..c8922b19dc5e 100644 --- a/tests/pipelines/chronoedit/test_chronoedit.py +++ b/tests/pipelines/chronoedit/test_chronoedit.py @@ -156,7 +156,7 @@ def test_inference(self): self.assertEqual(generated_video.shape, (5, 3, 16, 16)) # fmt: off - expected_slice = torch.tensor([0.4525, 0.4520, 0.4485, 0.4534, 0.4523, 0.4522, 0.4529, 0.4528, 0.5022, 0.5064, 0.5011, 0.5061, 0.5028, 0.4979, 0.5117, 0.5192]) + expected_slice = torch.tensor([0.4525, 0.4520, 0.4485, 0.4533, 0.4522, 0.4522, 0.4529, 0.4528, 0.5023, 0.5067, 0.5023, 0.5061, 0.5024, 0.4977, 0.5118, 0.5190]) # fmt: on generated_slice = generated_video.flatten() diff --git a/tests/pipelines/cogvideo/test_cogvideox_image2video.py b/tests/pipelines/cogvideo/test_cogvideox_image2video.py index c30cf56a933c..fc9e35c51b59 100644 --- a/tests/pipelines/cogvideo/test_cogvideox_image2video.py +++ b/tests/pipelines/cogvideo/test_cogvideox_image2video.py @@ -277,6 +277,7 @@ def test_vae_tiling(self, expected_diff_max: float = 0.3): # The reason to modify it this way is because I2V Transformer limits the generation to resolutions used during initialization. # This limitation comes from using learned positional embeddings which cannot be generated on-the-fly like sincos or RoPE embeddings. # See the if-statement on "self.use_learned_positional_embeddings" in diffusers/models/embeddings.py + torch.manual_seed(0) components["transformer"] = CogVideoXTransformer3DModel.from_config( components["transformer"].config, sample_height=16, diff --git a/tests/pipelines/controlnet/test_controlnet.py b/tests/pipelines/controlnet/test_controlnet.py index b142c2baf957..9b9df82141ea 100644 --- a/tests/pipelines/controlnet/test_controlnet.py +++ b/tests/pipelines/controlnet/test_controlnet.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -178,7 +178,7 @@ def test_attention_slicing_forward_pass(self): def test_ip_adapter(self): expected_pipe_slice = None if torch_device == "cpu": - expected_pipe_slice = np.array([0.5234, 0.3333, 0.1745, 0.7605, 0.6224, 0.4637, 0.6989, 0.7526, 0.4665]) + expected_pipe_slice = np.array([0.5245, 0.3353, 0.1784, 0.7556, 0.6239, 0.4673, 0.6931, 0.7487, 0.4644]) return super().test_ip_adapter(expected_pipe_slice=expected_pipe_slice) @unittest.skipIf( @@ -427,7 +427,7 @@ def test_inference_batch_single_identical(self): def test_ip_adapter(self): expected_pipe_slice = None if torch_device == "cpu": - expected_pipe_slice = np.array([0.2422, 0.3425, 0.4048, 0.5351, 0.3503, 0.2419, 0.4645, 0.4570, 0.3804]) + expected_pipe_slice = np.array([0.2395, 0.3421, 0.4023, 0.5345, 0.3496, 0.2402, 0.4645, 0.4563, 0.3786]) return super().test_ip_adapter(expected_pipe_slice=expected_pipe_slice) def test_save_pretrained_raise_not_implemented_exception(self): @@ -655,7 +655,7 @@ def test_inference_batch_single_identical(self): def test_ip_adapter(self): expected_pipe_slice = None if torch_device == "cpu": - expected_pipe_slice = np.array([0.5264, 0.3203, 0.1602, 0.8235, 0.6332, 0.4593, 0.7226, 0.7777, 0.4780]) + expected_pipe_slice = np.array([0.5234, 0.3198, 0.1596, 0.8201, 0.6316, 0.4566, 0.7209, 0.7761, 0.4757]) return super().test_ip_adapter(expected_pipe_slice=expected_pipe_slice) def test_save_pretrained_raise_not_implemented_exception(self): diff --git a/tests/pipelines/controlnet/test_controlnet_img2img.py b/tests/pipelines/controlnet/test_controlnet_img2img.py index c5d438e93427..d60b76190658 100644 --- a/tests/pipelines/controlnet/test_controlnet_img2img.py +++ b/tests/pipelines/controlnet/test_controlnet_img2img.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -177,7 +177,7 @@ def test_attention_slicing_forward_pass(self): def test_ip_adapter(self): expected_pipe_slice = None if torch_device == "cpu": - expected_pipe_slice = np.array([0.7096, 0.5149, 0.3571, 0.5897, 0.4715, 0.4052, 0.6098, 0.6886, 0.4213]) + expected_pipe_slice = np.array([0.7051, 0.5142, 0.3632, 0.5896, 0.4731, 0.4075, 0.6059, 0.6865, 0.4203]) return super().test_ip_adapter(expected_pipe_slice=expected_pipe_slice) @unittest.skipIf( @@ -384,7 +384,7 @@ def test_inference_batch_single_identical(self): def test_ip_adapter(self): expected_pipe_slice = None if torch_device == "cpu": - expected_pipe_slice = np.array([0.5293, 0.7339, 0.6642, 0.3950, 0.5212, 0.5175, 0.7002, 0.5907, 0.5182]) + expected_pipe_slice = np.array([0.5241, 0.7318, 0.6612, 0.3910, 0.5169, 0.5151, 0.6959, 0.5856, 0.5151]) return super().test_ip_adapter(expected_pipe_slice=expected_pipe_slice) def test_save_pretrained_raise_not_implemented_exception(self): diff --git a/tests/pipelines/controlnet/test_controlnet_inpaint.py b/tests/pipelines/controlnet/test_controlnet_inpaint.py index ebbe869e9e5e..c95f92dd8191 100644 --- a/tests/pipelines/controlnet/test_controlnet_inpaint.py +++ b/tests/pipelines/controlnet/test_controlnet_inpaint.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/pipelines/controlnet/test_controlnet_inpaint_sdxl.py b/tests/pipelines/controlnet/test_controlnet_inpaint_sdxl.py index c91f2c700c15..7f7976e839ed 100644 --- a/tests/pipelines/controlnet/test_controlnet_inpaint_sdxl.py +++ b/tests/pipelines/controlnet/test_controlnet_inpaint_sdxl.py @@ -343,7 +343,9 @@ def test_controlnet_sdxl_guess(self): output = sd_pipe(**inputs) image_slice = output.images[0, -3:, -3:, -1] - expected_slice = np.array([0.5460, 0.4943, 0.4635, 0.5832, 0.5366, 0.4815, 0.6034, 0.5741, 0.4341]) + expected_slice = np.array( + [0.559845, 0.506214, 0.467439, 0.587851, 0.536138, 0.480375, 0.598892, 0.571167, 0.434700] + ) # make sure that it's equal assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-4 diff --git a/tests/pipelines/controlnet/test_controlnet_sdxl.py b/tests/pipelines/controlnet/test_controlnet_sdxl.py index 42ec446dbfae..8736ee5fa060 100644 --- a/tests/pipelines/controlnet/test_controlnet_sdxl.py +++ b/tests/pipelines/controlnet/test_controlnet_sdxl.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -196,7 +196,7 @@ def test_ip_adapter(self, from_ssd1b=False, expected_pipe_slice=None): expected_pipe_slice = None if torch_device == "cpu": expected_pipe_slice = np.array( - [0.7335, 0.5866, 0.5623, 0.6242, 0.5751, 0.5999, 0.4091, 0.4590, 0.5054] + [0.7153, 0.5634, 0.5697, 0.6209, 0.5700, 0.6044, 0.4283, 0.4552, 0.5271] ) return super().test_ip_adapter(expected_pipe_slice=expected_pipe_slice) @@ -312,7 +312,9 @@ def test_controlnet_sdxl_guess(self): output = sd_pipe(**inputs) image_slice = output.images[0, -3:, -3:, -1] - expected_slice = np.array([0.7335, 0.5866, 0.5623, 0.6242, 0.5751, 0.5999, 0.4091, 0.4590, 0.5054]) + expected_slice = np.array( + [0.715316, 0.563373, 0.569716, 0.620860, 0.569999, 0.604369, 0.428272, 0.455195, 0.527119] + ) # make sure that it's equal assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-4 @@ -934,7 +936,9 @@ def test_controlnet_sdxl_guess(self): output = sd_pipe(**inputs) image_slice = output.images[0, -3:, -3:, -1] - expected_slice = np.array([0.7212, 0.5890, 0.5491, 0.6425, 0.5970, 0.6091, 0.4418, 0.4556, 0.5032]) + expected_slice = np.array( + [0.669912, 0.557802, 0.523260, 0.596366, 0.552897, 0.576922, 0.433411, 0.450273, 0.491615] + ) # make sure that it's equal assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-4 @@ -942,7 +946,7 @@ def test_controlnet_sdxl_guess(self): def test_ip_adapter(self): expected_pipe_slice = None if torch_device == "cpu": - expected_pipe_slice = np.array([0.7212, 0.5890, 0.5491, 0.6425, 0.5970, 0.6091, 0.4418, 0.4556, 0.5032]) + expected_pipe_slice = np.array([0.6699, 0.5578, 0.5233, 0.5964, 0.5529, 0.5769, 0.4334, 0.4503, 0.4916]) return super().test_ip_adapter(from_ssd1b=True, expected_pipe_slice=expected_pipe_slice) diff --git a/tests/pipelines/controlnet/test_controlnet_sdxl_img2img.py b/tests/pipelines/controlnet/test_controlnet_sdxl_img2img.py index bd4a233741e8..d242a1bacb92 100644 --- a/tests/pipelines/controlnet/test_controlnet_sdxl_img2img.py +++ b/tests/pipelines/controlnet/test_controlnet_sdxl_img2img.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -183,7 +183,7 @@ def get_dummy_inputs(self, device, seed=0): def test_ip_adapter(self): expected_pipe_slice = None if torch_device == "cpu": - expected_pipe_slice = np.array([0.6276, 0.5271, 0.5205, 0.5393, 0.5774, 0.5872, 0.5456, 0.5415, 0.5354]) + expected_pipe_slice = np.array([0.6710, 0.5497, 0.5469, 0.5758, 0.5990, 0.5996, 0.5583, 0.5506, 0.5368]) # TODO: update after slices.p return super().test_ip_adapter(expected_pipe_slice=expected_pipe_slice) @@ -200,7 +200,7 @@ def test_stable_diffusion_xl_controlnet_img2img(self): assert image.shape == (1, 64, 64, 3) expected_slice = np.array( - [0.5557202, 0.46418434, 0.46983826, 0.623529, 0.5557242, 0.49262643, 0.6070508, 0.5702978, 0.43777135] + [0.55813384, 0.4668495, 0.46676695, 0.6121852, 0.55514586, 0.49157068, 0.5960574, 0.56897247, 0.43931544] ) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 @@ -223,7 +223,7 @@ def test_stable_diffusion_xl_controlnet_img2img_guess(self): assert output.images.shape == (1, 64, 64, 3) expected_slice = np.array( - [0.5557202, 0.46418434, 0.46983826, 0.623529, 0.5557242, 0.49262643, 0.6070508, 0.5702978, 0.43777135] + [0.55813384, 0.4668495, 0.46676695, 0.6121852, 0.55514586, 0.49157068, 0.5960574, 0.56897247, 0.43931544] ) # make sure that it's equal diff --git a/tests/pipelines/controlnet_flux/test_controlnet_flux.py b/tests/pipelines/controlnet_flux/test_controlnet_flux.py index 8208a79904ed..6e223973eee0 100644 --- a/tests/pipelines/controlnet_flux/test_controlnet_flux.py +++ b/tests/pipelines/controlnet_flux/test_controlnet_flux.py @@ -174,7 +174,7 @@ def test_controlnet_flux(self): assert image.shape == (1, 32, 32, 3) - expected_slice = np.array([0.6677, 0.6138, 0.5296, 0.6109, 0.5672, 0.6373, 0.5463, 0.6068, 0.5569]) + expected_slice = np.array([0.6751, 0.6115, 0.5290, 0.6200, 0.5736, 0.6409, 0.5588, 0.6046, 0.5594]) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2, ( f"Expected: {expected_slice}, got: {image_slice.flatten()}" diff --git a/tests/pipelines/controlnet_sd3/test_controlnet_inpaint_sd3.py b/tests/pipelines/controlnet_sd3/test_controlnet_inpaint_sd3.py index 28abf122c41b..09bce003379b 100644 --- a/tests/pipelines/controlnet_sd3/test_controlnet_inpaint_sd3.py +++ b/tests/pipelines/controlnet_sd3/test_controlnet_inpaint_sd3.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -195,7 +195,7 @@ def test_controlnet_inpaint_sd3(self): assert image.shape == (1, 32, 32, 3) - expected_slice = np.array([0.2875, 0.3173, 0.4028, 0.7248, 0.6338, 0.4238, 0.1730, 0.4609, 0.5424]) + expected_slice = np.array([0.4627, 0.3686, 0.3741, 0.5855, 0.6071, 0.4046, 0.1916, 0.3938, 0.4953]) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2, ( f"Expected: {expected_slice}, got: {image_slice.flatten()}" diff --git a/tests/pipelines/controlnet_sd3/test_controlnet_sd3.py b/tests/pipelines/controlnet_sd3/test_controlnet_sd3.py index 4c8e80d38bfa..eca13af8340f 100644 --- a/tests/pipelines/controlnet_sd3/test_controlnet_sd3.py +++ b/tests/pipelines/controlnet_sd3/test_controlnet_sd3.py @@ -204,9 +204,9 @@ def run_pipe(self, components, use_sd35=False): assert image.shape == (1, 32, 32, 3) if not use_sd35: - expected_slice = np.array([0.4578, 0.3582, 0.4046, 0.0953, 0.6878, 0.5821, 0.5541, 0.5888, 0.4651]) + expected_slice = np.array([0.4121, 0.3775, 0.3734, 0.1509, 0.6324, 0.5503, 0.5425, 0.5614, 0.4061]) else: - expected_slice = np.array([0.3721, 0.5626, 0.4657, 0.2845, 0.5241, 0.5917, 0.6265, 0.6955, 0.3969]) + expected_slice = np.array([0.3793, 0.5179, 0.4389, 0.2820, 0.5148, 0.5565, 0.6282, 0.6891, 0.4197]) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2, ( f"Expected: {expected_slice}, got: {image_slice.flatten()}" diff --git a/tests/pipelines/ddim/test_ddim.py b/tests/pipelines/ddim/test_ddim.py index 731635bea605..1b5237f7726b 100644 --- a/tests/pipelines/ddim/test_ddim.py +++ b/tests/pipelines/ddim/test_ddim.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/pipelines/ddpm/test_ddpm.py b/tests/pipelines/ddpm/test_ddpm.py index 04ee741d8eb8..85676b56a1a3 100644 --- a/tests/pipelines/ddpm/test_ddpm.py +++ b/tests/pipelines/ddpm/test_ddpm.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/pipelines/deepfloyd_if/test_if.py b/tests/pipelines/deepfloyd_if/test_if.py index 0fd1391decd0..4fd555b9a33e 100644 --- a/tests/pipelines/deepfloyd_if/test_if.py +++ b/tests/pipelines/deepfloyd_if/test_if.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/pipelines/deepfloyd_if/test_if_img2img.py b/tests/pipelines/deepfloyd_if/test_if_img2img.py index 9d3c96052be6..d047f4a046b8 100644 --- a/tests/pipelines/deepfloyd_if/test_if_img2img.py +++ b/tests/pipelines/deepfloyd_if/test_if_img2img.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/pipelines/deepfloyd_if/test_if_img2img_superresolution.py b/tests/pipelines/deepfloyd_if/test_if_img2img_superresolution.py index e2114910edb0..d83341249ee0 100644 --- a/tests/pipelines/deepfloyd_if/test_if_img2img_superresolution.py +++ b/tests/pipelines/deepfloyd_if/test_if_img2img_superresolution.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/pipelines/deepfloyd_if/test_if_inpainting.py b/tests/pipelines/deepfloyd_if/test_if_inpainting.py index 2679e0b77690..915d400bc6e8 100644 --- a/tests/pipelines/deepfloyd_if/test_if_inpainting.py +++ b/tests/pipelines/deepfloyd_if/test_if_inpainting.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/pipelines/deepfloyd_if/test_if_inpainting_superresolution.py b/tests/pipelines/deepfloyd_if/test_if_inpainting_superresolution.py index 3d64556c6e41..142351385f40 100644 --- a/tests/pipelines/deepfloyd_if/test_if_inpainting_superresolution.py +++ b/tests/pipelines/deepfloyd_if/test_if_inpainting_superresolution.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/pipelines/deepfloyd_if/test_if_superresolution.py b/tests/pipelines/deepfloyd_if/test_if_superresolution.py index fa7c0fb2e062..7e0a9af1253c 100644 --- a/tests/pipelines/deepfloyd_if/test_if_superresolution.py +++ b/tests/pipelines/deepfloyd_if/test_if_superresolution.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/pipelines/dit/test_dit.py b/tests/pipelines/dit/test_dit.py index cd5c08ced3fc..406f3003c96a 100644 --- a/tests/pipelines/dit/test_dit.py +++ b/tests/pipelines/dit/test_dit.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/pipelines/hidream_image/test_pipeline_hidream.py b/tests/pipelines/hidream_image/test_pipeline_hidream.py index 1dcca2f5782d..fb41ebd21edd 100644 --- a/tests/pipelines/hidream_image/test_pipeline_hidream.py +++ b/tests/pipelines/hidream_image/test_pipeline_hidream.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -151,7 +151,7 @@ def test_inference(self): self.assertEqual(generated_image.shape, (128, 128, 3)) # fmt: off - expected_slice = np.array([0.4501, 0.5256, 0.4207, 0.5783, 0.4842, 0.4833, 0.4441, 0.5112, 0.6587, 0.3169, 0.7308, 0.5927, 0.6251, 0.5509, 0.5355, 0.5969]) + expected_slice = np.array([0.3841, 0.5881, 0.3958, 0.5585, 0.4646, 0.4385, 0.4529, 0.4178, 0.6595, 0.2871, 0.7331, 0.5203, 0.5525, 0.5296, 0.4695, 0.5524]) # fmt: on generated_slice = generated_image.flatten() diff --git a/tests/pipelines/hunyuan_video/test_hunyuan_image2video.py b/tests/pipelines/hunyuan_video/test_hunyuan_image2video.py index 4a0129e0826f..3b0f64e84c34 100644 --- a/tests/pipelines/hunyuan_video/test_hunyuan_image2video.py +++ b/tests/pipelines/hunyuan_video/test_hunyuan_image2video.py @@ -234,7 +234,7 @@ def test_inference(self): self.assertEqual(generated_video.shape, (5, 3, 16, 16)) # fmt: off - expected_slice = torch.tensor([0.4441, 0.4790, 0.4485, 0.5748, 0.3539, 0.1553, 0.2707, 0.3594, 0.5331, 0.6645, 0.6799, 0.5257, 0.5092, 0.3450, 0.4276, 0.4127]) + expected_slice = torch.tensor([0.4477, 0.4781, 0.4478, 0.5687, 0.3446, 0.1606, 0.2699, 0.3613, 0.5592, 0.6789, 0.6793, 0.5311, 0.5175, 0.3748, 0.4228, 0.4149]) # fmt: on generated_slice = generated_video.flatten() diff --git a/tests/pipelines/hunyuan_video/test_hunyuan_skyreels_image2video.py b/tests/pipelines/hunyuan_video/test_hunyuan_skyreels_image2video.py index 0ac15d4c405a..0102bfc64c3d 100644 --- a/tests/pipelines/hunyuan_video/test_hunyuan_skyreels_image2video.py +++ b/tests/pipelines/hunyuan_video/test_hunyuan_skyreels_image2video.py @@ -195,7 +195,7 @@ def test_inference(self): self.assertEqual(generated_video.shape, (9, 3, 16, 16)) # fmt: off - expected_slice = torch.tensor([0.5832, 0.5498, 0.4839, 0.4744, 0.4515, 0.4832, 0.496, 0.563, 0.5918, 0.5979, 0.5101, 0.6168, 0.6613, 0.536, 0.55, 0.5775]) + expected_slice = torch.tensor([0.5979, 0.5689, 0.5049, 0.4954, 0.4626, 0.5027, 0.4998, 0.5639, 0.5746, 0.5710, 0.5034, 0.5987, 0.6288, 0.5199, 0.5518, 0.5783]) # fmt: on generated_slice = generated_video.flatten() diff --git a/tests/pipelines/hunyuan_video/test_hunyuan_video.py b/tests/pipelines/hunyuan_video/test_hunyuan_video.py index 719bb9e67fbf..b16c8d7158a6 100644 --- a/tests/pipelines/hunyuan_video/test_hunyuan_video.py +++ b/tests/pipelines/hunyuan_video/test_hunyuan_video.py @@ -208,7 +208,7 @@ def test_inference(self): self.assertEqual(generated_video.shape, (9, 3, 16, 16)) # fmt: off - expected_slice = torch.tensor([0.3946, 0.4649, 0.3196, 0.4569, 0.3312, 0.3687, 0.3216, 0.3972, 0.4469, 0.3888, 0.3929, 0.3802, 0.3479, 0.3888, 0.3825, 0.3542]) + expected_slice = torch.tensor([0.3966, 0.4693, 0.3223, 0.4634, 0.3316, 0.3698, 0.3201, 0.3954, 0.4430, 0.3860, 0.3925, 0.3823, 0.3478, 0.3901, 0.3837, 0.3547]) # fmt: on generated_slice = generated_video.flatten() diff --git a/tests/pipelines/hunyuan_video/test_hunyuan_video_framepack.py b/tests/pipelines/hunyuan_video/test_hunyuan_video_framepack.py index f7f8908b6f8f..23fb17403f30 100644 --- a/tests/pipelines/hunyuan_video/test_hunyuan_video_framepack.py +++ b/tests/pipelines/hunyuan_video/test_hunyuan_video_framepack.py @@ -230,7 +230,7 @@ def test_inference(self): self.assertEqual(generated_video.shape, (13, 3, 32, 32)) # fmt: off - expected_slice = torch.tensor([0.363, 0.3384, 0.3426, 0.3512, 0.3372, 0.3276, 0.417, 0.4061, 0.5221, 0.467, 0.4813, 0.4556, 0.4107, 0.3945, 0.4049, 0.4551]) + expected_slice = torch.tensor([0.3628, 0.3380, 0.3421, 0.3505, 0.3362, 0.3268, 0.4167, 0.4063, 0.5225, 0.4693, 0.4827, 0.4583, 0.4144, 0.3983, 0.4089, 0.4587]) # fmt: on generated_slice = generated_video.flatten() diff --git a/tests/pipelines/hunyuandit/test_hunyuan_dit.py b/tests/pipelines/hunyuandit/test_hunyuan_dit.py index ba57b6a3599a..9b17ac606481 100644 --- a/tests/pipelines/hunyuandit/test_hunyuan_dit.py +++ b/tests/pipelines/hunyuandit/test_hunyuan_dit.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/pipelines/ip_adapters/test_ip_adapter_stable_diffusion.py b/tests/pipelines/ip_adapters/test_ip_adapter_stable_diffusion.py index 6689526e8b13..056521c7f55f 100644 --- a/tests/pipelines/ip_adapters/test_ip_adapter_stable_diffusion.py +++ b/tests/pipelines/ip_adapters/test_ip_adapter_stable_diffusion.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/pipelines/joyimage/test_joyimage_edit_plus.py b/tests/pipelines/joyimage/test_joyimage_edit_plus.py new file mode 100644 index 000000000000..e41265d30128 --- /dev/null +++ b/tests/pipelines/joyimage/test_joyimage_edit_plus.py @@ -0,0 +1,225 @@ +# Copyright 2025 The HuggingFace Team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import unittest +from unittest.mock import patch + +import numpy as np +import pytest +import torch +from PIL import Image +from transformers import Qwen3VLForConditionalGeneration, Qwen3VLProcessor + +from diffusers import ( + AutoencoderKLWan, + FlowMatchEulerDiscreteScheduler, + JoyImageEditPlusPipeline, + JoyImageEditPlusTransformer3DModel, +) +from diffusers.hooks import apply_group_offloading + +from ...testing_utils import enable_full_determinism, require_torch_accelerator, torch_device +from ..pipeline_params import TEXT_TO_IMAGE_PARAMS +from ..test_pipelines_common import PipelineTesterMixin + + +enable_full_determinism() + + +class JoyImageEditPlusPipelineFastTests(PipelineTesterMixin, unittest.TestCase): + pipeline_class = JoyImageEditPlusPipeline + params = TEXT_TO_IMAGE_PARAMS - {"cross_attention_kwargs"} + batch_params = frozenset(["prompt", "images"]) + required_optional_params = frozenset( + [ + "num_inference_steps", + "generator", + "latents", + "return_dict", + "callback_on_step_end", + "callback_on_step_end_tensor_inputs", + ] + ) + supports_dduf = False + test_xformers_attention = False + test_layerwise_casting = True + test_group_offloading = True + + def setUp(self): + super().setUp() + self._bucket_patcher = patch( + "diffusers.pipelines.joyimage.image_processor.find_best_bucket", + return_value=(32, 32), + ) + self._bucket_patcher.start() + + def tearDown(self): + self._bucket_patcher.stop() + super().tearDown() + + def get_dummy_components(self): + tiny_ckpt_id = "huangfeice/tiny-random-Qwen3VLForConditionalGeneration" + + torch.manual_seed(0) + transformer = JoyImageEditPlusTransformer3DModel( + patch_size=[1, 2, 2], + in_channels=16, + hidden_size=32, + num_attention_heads=2, + text_dim=16, + num_layers=1, + rope_dim_list=[4, 6, 6], + theta=256, + ) + + torch.manual_seed(0) + vae = AutoencoderKLWan( + base_dim=3, + z_dim=16, + dim_mult=[1, 1, 1, 1], + num_res_blocks=1, + temperal_downsample=[False, True, True], + ) + + scheduler = FlowMatchEulerDiscreteScheduler() + + processor = Qwen3VLProcessor.from_pretrained(tiny_ckpt_id) + processor.image_processor.min_pixels = 4 * 28 * 28 + processor.image_processor.max_pixels = 4 * 28 * 28 + + text_encoder = Qwen3VLForConditionalGeneration.from_pretrained(tiny_ckpt_id) + text_encoder.resize_token_embeddings(len(processor.tokenizer)) + + components = { + "transformer": transformer, + "vae": vae, + "scheduler": scheduler, + "text_encoder": text_encoder, + "tokenizer": processor.tokenizer, + "processor": processor, + } + return components + + def get_dummy_inputs(self, device, seed=0): + if str(device).startswith("mps"): + generator = torch.manual_seed(seed) + else: + generator = torch.Generator(device=device).manual_seed(seed) + + inputs = { + "prompt": "combine the two images", + "images": [Image.new("RGB", (32, 32)), Image.new("RGB", (32, 32))], + "generator": generator, + "num_inference_steps": 2, + "guidance_scale": 1.0, + "height": 32, + "width": 32, + "max_sequence_length": 16, + "output_type": "pt", + } + return inputs + + def test_inference(self): + device = "cpu" + + components = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe.to(device) + pipe.set_progress_bar_config(disable=None) + + inputs = self.get_dummy_inputs(device) + image = pipe(**inputs).images + generated_image = image[0] + + self.assertEqual(generated_image.shape, (3, 32, 32)) + + def test_inference_batch_single_identical(self): + self._test_inference_batch_single_identical(batch_size=3, expected_max_diff=1e-1) + + @unittest.skip("num_images_per_prompt not applicable: each prompt is bound to reference images") + def test_num_images_per_prompt(self): + pass + + @unittest.skip("Test not supported") + def test_attention_slicing_forward_pass(self): + pass + + @pytest.mark.xfail(condition=True, reason="Preconfigured embeddings need to be revisited.", strict=False) + def test_encode_prompt_works_in_isolation(self, extra_required_param_value_dict=None, atol=1e-4, rtol=1e-4): + super().test_encode_prompt_works_in_isolation(extra_required_param_value_dict, atol, rtol) + + @require_torch_accelerator + def test_group_offloading_inference(self): + if not self.test_group_offloading: + return + + def create_pipe(): + torch.manual_seed(0) + components = self.get_dummy_components() + pipe = self.pipeline_class(**components) + pipe.set_progress_bar_config(disable=None) + return pipe + + def run_forward(pipe): + torch.manual_seed(0) + inputs = self.get_dummy_inputs(torch_device) + return pipe(**inputs)[0] + + pipe = create_pipe().to(torch_device) + output_without_group_offloading = run_forward(pipe) + + pipe = create_pipe() + for component_name in ["transformer", "text_encoder"]: + component = getattr(pipe, component_name, None) + if component is None: + continue + if hasattr(component, "enable_group_offload"): + component.enable_group_offload( + torch.device(torch_device), offload_type="block_level", num_blocks_per_group=1 + ) + else: + apply_group_offloading( + component, + onload_device=torch.device(torch_device), + offload_type="block_level", + num_blocks_per_group=1, + ) + pipe.vae.to(torch_device) + output_with_block_level = run_forward(pipe) + + pipe = create_pipe() + pipe.transformer.enable_group_offload(torch.device(torch_device), offload_type="leaf_level") + pipe.text_encoder.to(torch_device) + pipe.vae.to(torch_device) + output_with_leaf_level = run_forward(pipe) + + if torch.is_tensor(output_without_group_offloading): + output_without_group_offloading = output_without_group_offloading.detach().cpu().numpy() + output_with_block_level = output_with_block_level.detach().cpu().numpy() + output_with_leaf_level = output_with_leaf_level.detach().cpu().numpy() + + self.assertTrue(np.allclose(output_without_group_offloading, output_with_block_level, atol=1e-4)) + self.assertTrue(np.allclose(output_without_group_offloading, output_with_leaf_level, atol=1e-4)) + + @unittest.skip("Qwen3VLForConditionalGeneration does not support leaf-level group offloading") + def test_pipeline_level_group_offloading_inference(self): + pass + + @unittest.skip("Qwen3VLForConditionalGeneration does not support sequential CPU offloading") + def test_sequential_cpu_offload_forward_pass(self): + pass + + @unittest.skip("Qwen3VLForConditionalGeneration does not support sequential CPU offloading") + def test_sequential_offload_forward_pass_twice(self): + pass diff --git a/tests/pipelines/kandinsky/test_kandinsky.py b/tests/pipelines/kandinsky/test_kandinsky.py index 6207e71df8cd..100dc510e0fb 100644 --- a/tests/pipelines/kandinsky/test_kandinsky.py +++ b/tests/pipelines/kandinsky/test_kandinsky.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/pipelines/kandinsky/test_kandinsky_combined.py b/tests/pipelines/kandinsky/test_kandinsky_combined.py index fdb36b433a94..fd63c55817ac 100644 --- a/tests/pipelines/kandinsky/test_kandinsky_combined.py +++ b/tests/pipelines/kandinsky/test_kandinsky_combined.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/pipelines/kandinsky/test_kandinsky_img2img.py b/tests/pipelines/kandinsky/test_kandinsky_img2img.py index 6d1b43a24fd9..df5034569b23 100644 --- a/tests/pipelines/kandinsky/test_kandinsky_img2img.py +++ b/tests/pipelines/kandinsky/test_kandinsky_img2img.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/pipelines/kandinsky/test_kandinsky_inpaint.py b/tests/pipelines/kandinsky/test_kandinsky_inpaint.py index e2f4aa2a4f14..c28efcf520bc 100644 --- a/tests/pipelines/kandinsky/test_kandinsky_inpaint.py +++ b/tests/pipelines/kandinsky/test_kandinsky_inpaint.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/pipelines/kandinsky/test_kandinsky_prior.py b/tests/pipelines/kandinsky/test_kandinsky_prior.py index 903a1e5decfa..ffccfffd86e3 100644 --- a/tests/pipelines/kandinsky/test_kandinsky_prior.py +++ b/tests/pipelines/kandinsky/test_kandinsky_prior.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -219,7 +219,7 @@ def test_kandinsky_prior(self): assert image.shape == (1, 32) expected_slice = np.array( - [-0.5948, 0.1875, -0.1523, -1.1995, -1.4061, -0.6367, -1.4607, -0.6406, 0.8793, -0.3891] + [-0.0171, 0.8655, -0.6831, 0.6393, -0.8142, -0.1628, -1.4405, -0.7309, 0.3505, -0.2847] ) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 diff --git a/tests/pipelines/kandinsky2_2/test_kandinsky.py b/tests/pipelines/kandinsky2_2/test_kandinsky.py index 38294aa4c111..293691ccdda8 100644 --- a/tests/pipelines/kandinsky2_2/test_kandinsky.py +++ b/tests/pipelines/kandinsky2_2/test_kandinsky.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/pipelines/kandinsky2_2/test_kandinsky_combined.py b/tests/pipelines/kandinsky2_2/test_kandinsky_combined.py index 30aa2896089b..7047fe2212bf 100644 --- a/tests/pipelines/kandinsky2_2/test_kandinsky_combined.py +++ b/tests/pipelines/kandinsky2_2/test_kandinsky_combined.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -94,7 +94,7 @@ def test_kandinsky(self): assert image.shape == (1, 64, 64, 3) - expected_slice = np.array([0.3076, 0.2729, 0.5668, 0.0522, 0.3384, 0.7028, 0.4908, 0.3659, 0.6243]) + expected_slice = np.array([0.1111, 0.0000, 0.6088, 0.2670, 0.3847, 0.8102, 0.4594, 0.4858, 0.5990]) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2, ( f" expected_slice {expected_slice}, but got {image_slice.flatten()}" @@ -224,7 +224,7 @@ def test_kandinsky(self): assert image.shape == (1, 64, 64, 3) - expected_slice = np.array([0.4445, 0.4287, 0.4596, 0.3919, 0.3730, 0.5039, 0.4834, 0.4269, 0.5521]) + expected_slice = np.array([0.4525, 0.4496, 0.4976, 0.4512, 0.4424, 0.5400, 0.4572, 0.4521, 0.5306]) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2, ( f" expected_slice {expected_slice}, but got {image_slice.flatten()}" diff --git a/tests/pipelines/kandinsky2_2/test_kandinsky_controlnet.py b/tests/pipelines/kandinsky2_2/test_kandinsky_controlnet.py index 8f8e58a8c4c8..e4b534bc1c4f 100644 --- a/tests/pipelines/kandinsky2_2/test_kandinsky_controlnet.py +++ b/tests/pipelines/kandinsky2_2/test_kandinsky_controlnet.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/pipelines/kandinsky2_2/test_kandinsky_controlnet_img2img.py b/tests/pipelines/kandinsky2_2/test_kandinsky_controlnet_img2img.py index a4346605929b..a45ac032f800 100644 --- a/tests/pipelines/kandinsky2_2/test_kandinsky_controlnet_img2img.py +++ b/tests/pipelines/kandinsky2_2/test_kandinsky_controlnet_img2img.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/pipelines/kandinsky2_2/test_kandinsky_img2img.py b/tests/pipelines/kandinsky2_2/test_kandinsky_img2img.py index 99f3fe0f40f1..9291854606bf 100644 --- a/tests/pipelines/kandinsky2_2/test_kandinsky_img2img.py +++ b/tests/pipelines/kandinsky2_2/test_kandinsky_img2img.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/pipelines/kandinsky2_2/test_kandinsky_inpaint.py b/tests/pipelines/kandinsky2_2/test_kandinsky_inpaint.py index 5659699fc7aa..5d19940ef6b3 100644 --- a/tests/pipelines/kandinsky2_2/test_kandinsky_inpaint.py +++ b/tests/pipelines/kandinsky2_2/test_kandinsky_inpaint.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/pipelines/kandinsky2_2/test_kandinsky_prior.py b/tests/pipelines/kandinsky2_2/test_kandinsky_prior.py index adcc6cc2167c..8775a413961e 100644 --- a/tests/pipelines/kandinsky2_2/test_kandinsky_prior.py +++ b/tests/pipelines/kandinsky2_2/test_kandinsky_prior.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -221,7 +221,7 @@ def test_kandinsky_prior(self): assert image.shape == (1, 32) expected_slice = np.array( - [-0.5948, 0.1875, -0.1523, -1.1995, -1.4061, -0.6367, -1.4607, -0.6406, 0.8793, -0.3891] + [-0.0171, 0.8655, -0.6831, 0.6393, -0.8142, -0.1628, -1.4405, -0.7309, 0.3505, -0.2847] ) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 diff --git a/tests/pipelines/kandinsky2_2/test_kandinsky_prior_emb2emb.py b/tests/pipelines/kandinsky2_2/test_kandinsky_prior_emb2emb.py index 5377d917791a..989676c52b6d 100644 --- a/tests/pipelines/kandinsky2_2/test_kandinsky_prior_emb2emb.py +++ b/tests/pipelines/kandinsky2_2/test_kandinsky_prior_emb2emb.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -223,7 +223,7 @@ def test_kandinsky_prior_emb2emb(self): assert image.shape == (1, 32) expected_slice = np.array( - [-0.8947, 0.7225, -0.2400, -1.4224, -1.9268, -1.1454, -1.8220, -0.7972, 1.0465, -0.5207] + [0.0117, 0.8621, -0.7459, 0.5970, -0.8612, -0.2034, -1.5705, -0.6786, 0.2857, -0.1696] ) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 diff --git a/tests/pipelines/kandinsky3/test_kandinsky3.py b/tests/pipelines/kandinsky3/test_kandinsky3.py index 7deae194fb4b..1a1aa4b9d9ca 100644 --- a/tests/pipelines/kandinsky3/test_kandinsky3.py +++ b/tests/pipelines/kandinsky3/test_kandinsky3.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/pipelines/kandinsky3/test_kandinsky3_img2img.py b/tests/pipelines/kandinsky3/test_kandinsky3_img2img.py index 469300e67788..b97aabe8eeaf 100644 --- a/tests/pipelines/kandinsky3/test_kandinsky3_img2img.py +++ b/tests/pipelines/kandinsky3/test_kandinsky3_img2img.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/pipelines/kolors/test_kolors.py b/tests/pipelines/kolors/test_kolors.py index f1d4982d4d74..7174e2a87d1f 100644 --- a/tests/pipelines/kolors/test_kolors.py +++ b/tests/pipelines/kolors/test_kolors.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/pipelines/kolors/test_kolors_img2img.py b/tests/pipelines/kolors/test_kolors_img2img.py index 5a5d31a46456..f0aec6c9af62 100644 --- a/tests/pipelines/kolors/test_kolors_img2img.py +++ b/tests/pipelines/kolors/test_kolors_img2img.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/pipelines/latent_consistency_models/test_latent_consistency_models.py b/tests/pipelines/latent_consistency_models/test_latent_consistency_models.py index c7666244b35f..c961fe97d1f6 100644 --- a/tests/pipelines/latent_consistency_models/test_latent_consistency_models.py +++ b/tests/pipelines/latent_consistency_models/test_latent_consistency_models.py @@ -112,7 +112,7 @@ def get_dummy_inputs(self, device, seed=0): def test_ip_adapter(self): expected_pipe_slice = None if torch_device == "cpu": - expected_pipe_slice = np.array([0.1403, 0.5072, 0.5316, 0.1202, 0.3865, 0.4211, 0.5363, 0.3557, 0.3645]) + expected_pipe_slice = np.array([0.1405, 0.5002, 0.5213, 0.1223, 0.3856, 0.4165, 0.5382, 0.3622, 0.3693]) return super().test_ip_adapter(expected_pipe_slice=expected_pipe_slice) def test_lcm_onestep(self): @@ -130,7 +130,7 @@ def test_lcm_onestep(self): assert image.shape == (1, 64, 64, 3) image_slice = image[0, -3:, -3:, -1] - expected_slice = np.array([0.1441, 0.5304, 0.5452, 0.1361, 0.4011, 0.4370, 0.5326, 0.3492, 0.3637]) + expected_slice = np.array([0.1444, 0.5229, 0.5344, 0.1384, 0.3999, 0.4320, 0.5345, 0.3559, 0.3685]) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3 def test_lcm_multistep(self): @@ -147,7 +147,7 @@ def test_lcm_multistep(self): assert image.shape == (1, 64, 64, 3) image_slice = image[0, -3:, -3:, -1] - expected_slice = np.array([0.1403, 0.5072, 0.5316, 0.1202, 0.3865, 0.4211, 0.5363, 0.3557, 0.3645]) + expected_slice = np.array([0.1405, 0.5002, 0.5213, 0.1223, 0.3856, 0.4165, 0.5382, 0.3622, 0.3693]) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3 def test_lcm_custom_timesteps(self): @@ -166,7 +166,7 @@ def test_lcm_custom_timesteps(self): assert image.shape == (1, 64, 64, 3) image_slice = image[0, -3:, -3:, -1] - expected_slice = np.array([0.1403, 0.5072, 0.5316, 0.1202, 0.3865, 0.4211, 0.5363, 0.3557, 0.3645]) + expected_slice = np.array([0.1405, 0.5002, 0.5213, 0.1223, 0.3856, 0.4165, 0.5382, 0.3622, 0.3693]) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3 def test_inference_batch_single_identical(self): diff --git a/tests/pipelines/latent_consistency_models/test_latent_consistency_models_img2img.py b/tests/pipelines/latent_consistency_models/test_latent_consistency_models_img2img.py index d8e7745b7805..d245ebafcd9f 100644 --- a/tests/pipelines/latent_consistency_models/test_latent_consistency_models_img2img.py +++ b/tests/pipelines/latent_consistency_models/test_latent_consistency_models_img2img.py @@ -123,7 +123,7 @@ def get_dummy_inputs(self, device, seed=0): def test_ip_adapter(self): expected_pipe_slice = None if torch_device == "cpu": - expected_pipe_slice = np.array([0.4003, 0.3718, 0.2863, 0.5500, 0.5587, 0.3772, 0.4617, 0.4961, 0.4417]) + expected_pipe_slice = np.array([0.3952, 0.3674, 0.2839, 0.5461, 0.5532, 0.3737, 0.4652, 0.4986, 0.4422]) return super().test_ip_adapter(expected_pipe_slice=expected_pipe_slice) def test_lcm_onestep(self): @@ -141,7 +141,7 @@ def test_lcm_onestep(self): assert image.shape == (1, 32, 32, 3) image_slice = image[0, -3:, -3:, -1] - expected_slice = np.array([0.4388, 0.3717, 0.2202, 0.7213, 0.6370, 0.3664, 0.5815, 0.6080, 0.4977]) + expected_slice = np.array([0.4317, 0.3653, 0.2190, 0.7136, 0.6321, 0.3634, 0.5846, 0.6095, 0.4969]) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3 def test_lcm_multistep(self): @@ -158,7 +158,7 @@ def test_lcm_multistep(self): assert image.shape == (1, 32, 32, 3) image_slice = image[0, -3:, -3:, -1] - expected_slice = np.array([0.4150, 0.3719, 0.2479, 0.6333, 0.6024, 0.3778, 0.5036, 0.5420, 0.4678]) + expected_slice = np.array([0.4083, 0.3668, 0.2467, 0.6268, 0.5976, 0.3750, 0.5071, 0.5439, 0.4677]) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3 def test_lcm_custom_timesteps(self): @@ -177,7 +177,7 @@ def test_lcm_custom_timesteps(self): assert image.shape == (1, 32, 32, 3) image_slice = image[0, -3:, -3:, -1] - expected_slice = np.array([0.3994, 0.3471, 0.2540, 0.7030, 0.6193, 0.3645, 0.5777, 0.5850, 0.4965]) + expected_slice = np.array([0.3985, 0.3444, 0.2534, 0.6969, 0.6167, 0.3622, 0.5754, 0.5850, 0.4959]) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3 def test_inference_batch_single_identical(self): diff --git a/tests/pipelines/latent_diffusion/test_latent_diffusion.py b/tests/pipelines/latent_diffusion/test_latent_diffusion.py index 21c5bcf5a5b9..e87ea99a3805 100644 --- a/tests/pipelines/latent_diffusion/test_latent_diffusion.py +++ b/tests/pipelines/latent_diffusion/test_latent_diffusion.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -131,7 +131,7 @@ def test_inference_text2img(self): image_slice = image[0, -3:, -3:, -1] assert image.shape == (1, 16, 16, 3) - expected_slice = np.array([0.6101, 0.6156, 0.5622, 0.4895, 0.6661, 0.3804, 0.5748, 0.6136, 0.5014]) + expected_slice = np.array([0.6511, 0.5873, 0.5183, 0.5046, 0.6663, 0.3908, 0.5729, 0.5979, 0.5058]) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3 diff --git a/tests/pipelines/latent_diffusion/test_latent_diffusion_superresolution.py b/tests/pipelines/latent_diffusion/test_latent_diffusion_superresolution.py index b2cbdb9f5b45..0a39c787a236 100644 --- a/tests/pipelines/latent_diffusion/test_latent_diffusion_superresolution.py +++ b/tests/pipelines/latent_diffusion/test_latent_diffusion_superresolution.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/pipelines/marigold/test_marigold_depth.py b/tests/pipelines/marigold/test_marigold_depth.py index 04c55e7dde75..266ccc17cdd9 100644 --- a/tests/pipelines/marigold/test_marigold_depth.py +++ b/tests/pipelines/marigold/test_marigold_depth.py @@ -178,13 +178,13 @@ def _test_marigold_depth( def test_marigold_depth_dummy_defaults(self): self._test_marigold_depth( - expected_slice=np.array([0.4529, 0.5184, 0.4985, 0.4355, 0.4273, 0.4153, 0.5229, 0.4818, 0.4627]), + expected_slice=np.array([0.43442, 0.51455, 0.48409, 0.43800, 0.43542, 0.41082, 0.52997, 0.48687, 0.45823]), ) def test_marigold_depth_dummy_G0_S1_P32_E1_B1_M1(self): self._test_marigold_depth( generator_seed=0, - expected_slice=np.array([0.4529, 0.5184, 0.4985, 0.4355, 0.4273, 0.4153, 0.5229, 0.4818, 0.4627]), + expected_slice=np.array([0.43442, 0.51455, 0.48409, 0.43800, 0.43542, 0.41082, 0.52997, 0.48687, 0.45823]), num_inference_steps=1, processing_resolution=32, ensemble_size=1, @@ -195,7 +195,7 @@ def test_marigold_depth_dummy_G0_S1_P32_E1_B1_M1(self): def test_marigold_depth_dummy_G0_S1_P16_E1_B1_M1(self): self._test_marigold_depth( generator_seed=0, - expected_slice=np.array([0.4511, 0.4531, 0.4542, 0.5024, 0.4987, 0.4969, 0.5281, 0.5215, 0.5182]), + expected_slice=np.array([0.44393, 0.46028, 0.46846, 0.49471, 0.49320, 0.49244, 0.52010, 0.50965, 0.50443]), num_inference_steps=1, processing_resolution=16, ensemble_size=1, @@ -206,7 +206,7 @@ def test_marigold_depth_dummy_G0_S1_P16_E1_B1_M1(self): def test_marigold_depth_dummy_G2024_S1_P32_E1_B1_M1(self): self._test_marigold_depth( generator_seed=2024, - expected_slice=np.array([0.4671, 0.4739, 0.5130, 0.4308, 0.4411, 0.4720, 0.5064, 0.4796, 0.4795]), + expected_slice=np.array([0.48864, 0.47408, 0.51305, 0.43479, 0.43492, 0.46720, 0.50389, 0.48094, 0.47948]), num_inference_steps=1, processing_resolution=32, ensemble_size=1, @@ -217,7 +217,7 @@ def test_marigold_depth_dummy_G2024_S1_P32_E1_B1_M1(self): def test_marigold_depth_dummy_G0_S2_P32_E1_B1_M1(self): self._test_marigold_depth( generator_seed=0, - expected_slice=np.array([0.4165, 0.4485, 0.4647, 0.4003, 0.4577, 0.5074, 0.5106, 0.5077, 0.5042]), + expected_slice=np.array([0.40830, 0.45729, 0.46504, 0.39601, 0.45839, 0.51121, 0.51142, 0.50824, 0.50636]), num_inference_steps=2, processing_resolution=32, ensemble_size=1, @@ -228,7 +228,7 @@ def test_marigold_depth_dummy_G0_S2_P32_E1_B1_M1(self): def test_marigold_depth_dummy_G0_S1_P64_E1_B1_M1(self): self._test_marigold_depth( generator_seed=0, - expected_slice=np.array([0.4817, 0.5425, 0.5146, 0.5367, 0.5034, 0.4743, 0.4395, 0.4734, 0.4399]), + expected_slice=np.array([0.47847, 0.53579, 0.50407, 0.54443, 0.50714, 0.47101, 0.44327, 0.46812, 0.43958]), num_inference_steps=1, processing_resolution=64, ensemble_size=1, @@ -265,7 +265,7 @@ def test_marigold_depth_dummy_G0_S1_P32_E4_B2_M1(self): def test_marigold_depth_dummy_G0_S1_P16_E1_B1_M0(self): self._test_marigold_depth( generator_seed=0, - expected_slice=np.array([0.5515, 0.4588, 0.4197, 0.4741, 0.4229, 0.4328, 0.5333, 0.5314, 0.5182]), + expected_slice=np.array([0.53228, 0.46153, 0.42818, 0.46746, 0.40590, 0.45647, 0.52804, 0.52532, 0.50443]), num_inference_steps=1, processing_resolution=16, ensemble_size=1, diff --git a/tests/pipelines/marigold/test_marigold_intrinsics.py b/tests/pipelines/marigold/test_marigold_intrinsics.py index 6bed83a4001a..73da307b24fc 100644 --- a/tests/pipelines/marigold/test_marigold_intrinsics.py +++ b/tests/pipelines/marigold/test_marigold_intrinsics.py @@ -314,7 +314,7 @@ def test_marigold_depth_dummy_G0_S1_P16_E1_B1_M1(self): def test_marigold_depth_dummy_G2024_S1_P32_E1_B1_M1(self): self._test_marigold_intrinsics( generator_seed=2024, - expected_slice=np.array([0.40257, 0.39468, 0.51373, 0.4161, 0.40162, 0.58535, 0.43581, 0.47834, 0.48951]), + expected_slice=np.array([0.40250, 0.39464, 0.51378, 0.41603, 0.40150, 0.58531, 0.43581, 0.47833, 0.48946]), num_inference_steps=1, processing_resolution=32, ensemble_size=1, @@ -325,7 +325,7 @@ def test_marigold_depth_dummy_G2024_S1_P32_E1_B1_M1(self): def test_marigold_depth_dummy_G0_S2_P32_E1_B1_M1(self): self._test_marigold_intrinsics( generator_seed=0, - expected_slice=np.array([0.49636, 0.4518, 0.42722, 0.59044, 0.6362, 0.39011, 0.53522, 0.55153, 0.48699]), + expected_slice=np.array([0.52018, 0.45545, 0.42104, 0.58673, 0.63164, 0.38469, 0.52228, 0.54939, 0.48622]), num_inference_steps=2, processing_resolution=32, ensemble_size=1, @@ -336,7 +336,7 @@ def test_marigold_depth_dummy_G0_S2_P32_E1_B1_M1(self): def test_marigold_depth_dummy_G0_S1_P64_E1_B1_M1(self): self._test_marigold_intrinsics( generator_seed=0, - expected_slice=np.array([0.55547, 0.43511, 0.4887, 0.56399, 0.63867, 0.56337, 0.47889, 0.52925, 0.49235]), + expected_slice=np.array([0.55574, 0.43518, 0.48871, 0.56418, 0.63882, 0.56345, 0.47897, 0.52932, 0.49240]), num_inference_steps=1, processing_resolution=64, ensemble_size=1, @@ -347,7 +347,7 @@ def test_marigold_depth_dummy_G0_S1_P64_E1_B1_M1(self): def test_marigold_depth_dummy_G0_S1_P32_E3_B1_M1(self): self._test_marigold_intrinsics( generator_seed=0, - expected_slice=np.array([0.57249, 0.49824, 0.54438, 0.57733, 0.52404, 0.5255, 0.56493, 0.56336, 0.48579]), + expected_slice=np.array([0.57244, 0.49813, 0.54442, 0.57727, 0.52388, 0.52545, 0.56492, 0.56334, 0.48579]), num_inference_steps=1, processing_resolution=32, ensemble_size=3, @@ -359,7 +359,7 @@ def test_marigold_depth_dummy_G0_S1_P32_E3_B1_M1(self): def test_marigold_depth_dummy_G0_S1_P32_E4_B2_M1(self): self._test_marigold_intrinsics( generator_seed=0, - expected_slice=np.array([0.6294, 0.5575, 0.53414, 0.61077, 0.57156, 0.53974, 0.52956, 0.55467, 0.48751]), + expected_slice=np.array([0.62939, 0.55744, 0.53417, 0.61068, 0.57141, 0.53967, 0.52955, 0.55467, 0.48751]), num_inference_steps=1, processing_resolution=32, ensemble_size=4, @@ -371,7 +371,7 @@ def test_marigold_depth_dummy_G0_S1_P32_E4_B2_M1(self): def test_marigold_depth_dummy_G0_S1_P16_E1_B1_M0(self): self._test_marigold_intrinsics( generator_seed=0, - expected_slice=np.array([0.63511, 0.68137, 0.48783, 0.46689, 0.58505, 0.36757, 0.58465, 0.54302, 0.50387]), + expected_slice=np.array([0.63543, 0.68147, 0.48780, 0.46715, 0.58511, 0.36761, 0.58482, 0.54309, 0.50388]), num_inference_steps=1, processing_resolution=16, ensemble_size=1, diff --git a/tests/pipelines/marigold/test_marigold_normals.py b/tests/pipelines/marigold/test_marigold_normals.py index 108163bf22ec..e9f9ac7121c6 100644 --- a/tests/pipelines/marigold/test_marigold_normals.py +++ b/tests/pipelines/marigold/test_marigold_normals.py @@ -176,13 +176,17 @@ def _test_marigold_normals( def test_marigold_depth_dummy_defaults(self): self._test_marigold_normals( - expected_slice=np.array([0.0967, 0.5234, 0.1448, -0.3155, -0.2550, -0.5578, 0.6854, 0.5657, -0.1263]), + expected_slice=np.array( + [0.01655, 0.54110, 0.01681, -0.27346, -0.16697, -0.55219, 0.63358, 0.57275, -0.26173] + ), ) def test_marigold_depth_dummy_G0_S1_P32_E1_B1_M1(self): self._test_marigold_normals( generator_seed=0, - expected_slice=np.array([0.0967, 0.5234, 0.1448, -0.3155, -0.2550, -0.5578, 0.6854, 0.5657, -0.1263]), + expected_slice=np.array( + [0.01655, 0.54110, 0.01681, -0.27346, -0.16697, -0.55219, 0.63358, 0.57275, -0.26173] + ), num_inference_steps=1, processing_resolution=32, ensemble_size=1, @@ -193,7 +197,9 @@ def test_marigold_depth_dummy_G0_S1_P32_E1_B1_M1(self): def test_marigold_depth_dummy_G0_S1_P16_E1_B1_M1(self): self._test_marigold_normals( generator_seed=0, - expected_slice=np.array([-0.4128, -0.5918, -0.6540, 0.2446, -0.2687, -0.4607, 0.2935, -0.0483, -0.2086]), + expected_slice=np.array( + [-0.46928, -0.23894, -0.10984, -0.27850, -0.53089, -0.58686, -0.09792, -0.36364, -0.46909] + ), num_inference_steps=1, processing_resolution=16, ensemble_size=1, @@ -204,7 +210,9 @@ def test_marigold_depth_dummy_G0_S1_P16_E1_B1_M1(self): def test_marigold_depth_dummy_G2024_S1_P32_E1_B1_M1(self): self._test_marigold_normals( generator_seed=2024, - expected_slice=np.array([0.5731, -0.7631, -0.0199, 0.1609, -0.4628, -0.7044, 0.5761, -0.3471, -0.4498]), + expected_slice=np.array( + [0.75023, -0.90784, -0.08686, 0.07177, -0.59057, -0.73950, 0.52375, -0.26714, -0.43062] + ), num_inference_steps=1, processing_resolution=32, ensemble_size=1, @@ -215,7 +223,9 @@ def test_marigold_depth_dummy_G2024_S1_P32_E1_B1_M1(self): def test_marigold_depth_dummy_G0_S2_P32_E1_B1_M1(self): self._test_marigold_normals( generator_seed=0, - expected_slice=np.array([0.1017, -0.6823, -0.2533, 0.1988, 0.3389, 0.8478, 0.7757, 0.5220, 0.8668]), + expected_slice=np.array( + [0.04868, -0.58444, -0.26729, 0.13351, 0.36448, 0.86063, 0.74093, 0.58727, 0.91568] + ), num_inference_steps=2, processing_resolution=32, ensemble_size=1, @@ -226,7 +236,9 @@ def test_marigold_depth_dummy_G0_S2_P32_E1_B1_M1(self): def test_marigold_depth_dummy_G0_S1_P64_E1_B1_M1(self): self._test_marigold_normals( generator_seed=0, - expected_slice=np.array([-0.2391, 0.7969, 0.6224, 0.0698, 0.5669, -0.2167, -0.1362, -0.8945, -0.5501]), + expected_slice=np.array( + [-0.31085, 0.85586, 0.43578, 0.23959, 0.64753, -0.33613, -0.02879, -0.78712, -0.56993] + ), num_inference_steps=1, processing_resolution=64, ensemble_size=1, @@ -237,7 +249,9 @@ def test_marigold_depth_dummy_G0_S1_P64_E1_B1_M1(self): def test_marigold_depth_dummy_G0_S1_P32_E3_B1_M1(self): self._test_marigold_normals( generator_seed=0, - expected_slice=np.array([0.3826, -0.9634, -0.3835, 0.3514, 0.0691, -0.6182, 0.8709, 0.1590, -0.2181]), + expected_slice=np.array( + [0.28313, -0.91824, -0.38751, 0.35233, 0.13069, -0.58350, 0.82024, 0.04305, -0.27604] + ), num_inference_steps=1, processing_resolution=32, ensemble_size=3, @@ -249,7 +263,9 @@ def test_marigold_depth_dummy_G0_S1_P32_E3_B1_M1(self): def test_marigold_depth_dummy_G0_S1_P32_E4_B2_M1(self): self._test_marigold_normals( generator_seed=0, - expected_slice=np.array([0.2500, -0.3928, -0.2415, 0.1133, 0.2357, -0.4223, 0.9967, 0.4859, -0.1282]), + expected_slice=np.array( + [0.09912, -0.29170, -0.23348, 0.14560, 0.27345, -0.40117, 0.98967, 0.38735, -0.21680] + ), num_inference_steps=1, processing_resolution=32, ensemble_size=4, @@ -261,7 +277,9 @@ def test_marigold_depth_dummy_G0_S1_P32_E4_B2_M1(self): def test_marigold_depth_dummy_G0_S1_P16_E1_B1_M0(self): self._test_marigold_normals( generator_seed=0, - expected_slice=np.array([0.9588, 0.3326, -0.0825, -0.0994, -0.3534, -0.4302, 0.3562, 0.4421, -0.2086]), + expected_slice=np.array( + [0.87316, 0.43462, -0.15760, 0.20502, -0.42155, 0.07312, 0.36621, 0.03301, -0.46909] + ), num_inference_steps=1, processing_resolution=16, ensemble_size=1, diff --git a/tests/pipelines/ovis_image/test_ovis_image.py b/tests/pipelines/ovis_image/test_ovis_image.py index ac1670535363..6570d180454e 100644 --- a/tests/pipelines/ovis_image/test_ovis_image.py +++ b/tests/pipelines/ovis_image/test_ovis_image.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/pipelines/pag/test_pag_animatediff.py b/tests/pipelines/pag/test_pag_animatediff.py index b1cbd82d7679..ba508c88a277 100644 --- a/tests/pipelines/pag/test_pag_animatediff.py +++ b/tests/pipelines/pag/test_pag_animatediff.py @@ -181,33 +181,33 @@ def test_ip_adapter(self): if torch_device == "cpu": expected_pipe_slice = np.array( [ - 0.5068, - 0.5294, - 0.4926, - 0.4810, - 0.4188, - 0.5935, - 0.5295, - 0.3947, - 0.5300, - 0.4706, - 0.3950, - 0.4737, - 0.4072, - 0.3227, - 0.5481, - 0.4864, - 0.4518, - 0.5315, - 0.5979, - 0.5374, - 0.3503, - 0.5275, - 0.6067, - 0.4914, - 0.5440, - 0.4775, - 0.5538, + 0.5254, + 0.5844, + 0.4705, + 0.4952, + 0.4953, + 0.5932, + 0.5227, + 0.4285, + 0.5430, + 0.4903, + 0.4027, + 0.4942, + 0.3771, + 0.3899, + 0.6014, + 0.4469, + 0.4935, + 0.5543, + 0.5867, + 0.5314, + 0.3727, + 0.5183, + 0.6138, + 0.5170, + 0.5111, + 0.4725, + 0.5741, ] ) return super().test_ip_adapter(expected_pipe_slice=expected_pipe_slice) @@ -215,7 +215,7 @@ def test_ip_adapter(self): def test_dict_tuple_outputs_equivalent(self): expected_slice = None if torch_device == "cpu": - expected_slice = np.array([0.5295, 0.3947, 0.5300, 0.4864, 0.4518, 0.5315, 0.5440, 0.4775, 0.5538]) + expected_slice = np.array([0.5227, 0.4285, 0.5430, 0.4469, 0.4935, 0.5543, 0.5111, 0.4725, 0.5741]) return super().test_dict_tuple_outputs_equivalent(expected_slice=expected_slice) @require_accelerator diff --git a/tests/pipelines/pag/test_pag_controlnet_sd.py b/tests/pipelines/pag/test_pag_controlnet_sd.py index 36d5ae100a58..c65fadff3bd9 100644 --- a/tests/pipelines/pag/test_pag_controlnet_sd.py +++ b/tests/pipelines/pag/test_pag_controlnet_sd.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -213,7 +213,7 @@ def test_pag_cfg(self): 3, ), f"the shape of the output image should be (1, 64, 64, 3) but got {image.shape}" expected_slice = np.array( - [0.45505235, 0.2785938, 0.16334778, 0.79689944, 0.53095645, 0.40135607, 0.7052706, 0.69065094, 0.41548574] + [0.4533307, 0.28746086, 0.16543725, 0.786471, 0.5409842, 0.40167668, 0.7050086, 0.69445, 0.42238155] ) max_diff = np.abs(image_slice.flatten() - expected_slice).max() @@ -239,7 +239,7 @@ def test_pag_uncond(self): 3, ), f"the shape of the output image should be (1, 64, 64, 3) but got {image.shape}" expected_slice = np.array( - [0.45127502, 0.2797252, 0.15970308, 0.7993157, 0.5414344, 0.40160775, 0.7114598, 0.69803864, 0.4217583] + [0.45054412, 0.27958393, 0.15983358, 0.80098593, 0.5432001, 0.4018666, 0.7127423, 0.69932824, 0.42260933] ) max_diff = np.abs(image_slice.flatten() - expected_slice).max() diff --git a/tests/pipelines/pag/test_pag_controlnet_sd_inpaint.py b/tests/pipelines/pag/test_pag_controlnet_sd_inpaint.py index 948381f9769e..a59180cb9540 100644 --- a/tests/pipelines/pag/test_pag_controlnet_sd_inpaint.py +++ b/tests/pipelines/pag/test_pag_controlnet_sd_inpaint.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -209,7 +209,7 @@ def test_pag_cfg(self): 3, ), f"the shape of the output image should be (1, 64, 64, 3) but got {image.shape}" expected_slice = np.array( - [0.7488756, 0.61194265, 0.53382546, 0.5993959, 0.6193306, 0.56880975, 0.41277143, 0.5050145, 0.49376273] + [0.7277897, 0.61666954, 0.54722667, 0.595576, 0.593909, 0.56389576, 0.41761285, 0.50566983, 0.49766505] ) max_diff = np.abs(image_slice.flatten() - expected_slice).max() @@ -235,7 +235,7 @@ def test_pag_uncond(self): 3, ), f"the shape of the output image should be (1, 64, 64, 3) but got {image.shape}" expected_slice = np.array( - [0.7410303, 0.5989337, 0.530866, 0.60571927, 0.6162597, 0.5719856, 0.4187478, 0.5101238, 0.4978468] + [0.7349223, 0.60567534, 0.5428778, 0.6091342, 0.60273147, 0.57611704, 0.42401767, 0.5064247, 0.49535546] ) max_diff = np.abs(image_slice.flatten() - expected_slice).max() diff --git a/tests/pipelines/pag/test_pag_controlnet_sdxl.py b/tests/pipelines/pag/test_pag_controlnet_sdxl.py index 51b00f6932bc..25d0c77d8627 100644 --- a/tests/pipelines/pag/test_pag_controlnet_sdxl.py +++ b/tests/pipelines/pag/test_pag_controlnet_sdxl.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -230,7 +230,7 @@ def test_pag_cfg(self): 64, 3, ), f"the shape of the output image should be (1, 64, 64, 3) but got {image.shape}" - expected_slice = np.array([0.7036, 0.5613, 0.5526, 0.6129, 0.5610, 0.5842, 0.4228, 0.4612, 0.5017]) + expected_slice = np.array([0.6864, 0.5436, 0.5644, 0.6136, 0.5541, 0.5910, 0.4519, 0.4634, 0.5252]) max_diff = np.abs(image_slice.flatten() - expected_slice).max() assert max_diff < 1e-3, f"output is different from expected, {image_slice.flatten()}" @@ -254,7 +254,7 @@ def test_pag_uncond(self): 64, 3, ), f"the shape of the output image should be (1, 64, 64, 3) but got {image.shape}" - expected_slice = np.array([0.6888, 0.5398, 0.5603, 0.6086, 0.5541, 0.5957, 0.4332, 0.4643, 0.5154]) + expected_slice = np.array([0.6843, 0.5381, 0.5675, 0.6109, 0.5493, 0.5988, 0.4477, 0.4679, 0.5242]) max_diff = np.abs(image_slice.flatten() - expected_slice).max() assert max_diff < 1e-3, f"output is different from expected, {image_slice.flatten()}" diff --git a/tests/pipelines/pag/test_pag_controlnet_sdxl_img2img.py b/tests/pipelines/pag/test_pag_controlnet_sdxl_img2img.py index 3c1088adbcf2..e23f92e31119 100644 --- a/tests/pipelines/pag/test_pag_controlnet_sdxl_img2img.py +++ b/tests/pipelines/pag/test_pag_controlnet_sdxl_img2img.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -236,7 +236,7 @@ def test_pag_cfg(self): 3, ), f"the shape of the output image should be (1, 64, 64, 3) but got {image.shape}" expected_slice = np.array( - [0.5562928, 0.44882968, 0.4588066, 0.63200223, 0.5694165, 0.4955688, 0.6126959, 0.57588536, 0.43827885] + [0.55155665, 0.4650753, 0.46541628, 0.60965055, 0.55995595, 0.49751496, 0.5937391, 0.5700847, 0.44238678] ) max_diff = np.abs(image_slice.flatten() - expected_slice).max() @@ -262,7 +262,7 @@ def test_pag_uncond(self): 3, ), f"the shape of the output image should be (1, 64, 64, 3) but got {image.shape}" expected_slice = np.array( - [0.5543988, 0.45614323, 0.4665692, 0.6202247, 0.5598917, 0.49621183, 0.6084159, 0.5722314, 0.43945464] + [0.549061, 0.46218234, 0.4675981, 0.6109464, 0.5547849, 0.4960261, 0.60211027, 0.5698843, 0.44092298] ) max_diff = np.abs(image_slice.flatten() - expected_slice).max() diff --git a/tests/pipelines/pag/test_pag_hunyuan_dit.py b/tests/pipelines/pag/test_pag_hunyuan_dit.py index 38686ee448de..d4dd144a840b 100644 --- a/tests/pipelines/pag/test_pag_hunyuan_dit.py +++ b/tests/pipelines/pag/test_pag_hunyuan_dit.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/pipelines/pag/test_pag_kolors.py b/tests/pipelines/pag/test_pag_kolors.py index 1bbb4e79e4bc..a2cd926cfd56 100644 --- a/tests/pipelines/pag/test_pag_kolors.py +++ b/tests/pipelines/pag/test_pag_kolors.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/pipelines/pag/test_pag_pixart_sigma.py b/tests/pipelines/pag/test_pag_pixart_sigma.py index 9bc2f6eed395..2c0b89d0cfe0 100644 --- a/tests/pipelines/pag/test_pag_pixart_sigma.py +++ b/tests/pipelines/pag/test_pag_pixart_sigma.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/pipelines/pag/test_pag_sd.py b/tests/pipelines/pag/test_pag_sd.py index 064815d13693..1dd3ef298fd0 100644 --- a/tests/pipelines/pag/test_pag_sd.py +++ b/tests/pipelines/pag/test_pag_sd.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -272,7 +272,7 @@ def test_pag_inference(self): ), f"the shape of the output image should be (1, 64, 64, 3) but got {image.shape}" expected_slice = np.array( - [0.22802538, 0.44626093, 0.48905736, 0.29633686, 0.36400637, 0.4724258, 0.4678891, 0.32260418, 0.41611585] + [0.23171297, 0.44669262, 0.48407662, 0.29981518, 0.36721927, 0.46788025, 0.46333545, 0.3314417, 0.42078137] ) max_diff = np.abs(image_slice.flatten() - expected_slice).max() self.assertLessEqual(max_diff, 1e-3) diff --git a/tests/pipelines/pag/test_pag_sd3_img2img.py b/tests/pipelines/pag/test_pag_sd3_img2img.py index e4146b87803c..ede146915c55 100644 --- a/tests/pipelines/pag/test_pag_sd3_img2img.py +++ b/tests/pipelines/pag/test_pag_sd3_img2img.py @@ -196,7 +196,7 @@ def test_pag_inference(self): ), f"the shape of the output image should be (1, 32, 32, 3) but got {image.shape}" expected_slice = np.array( - [0.66063476, 0.44838923, 0.5484299, 0.7242875, 0.5970012, 0.6015729, 0.53080845, 0.52220416, 0.56397927] + [0.741577, 0.5491905, 0.59911674, 0.7702221, 0.65531653, 0.60989463, 0.491042, 0.5380116, 0.5592475] ) max_diff = np.abs(image_slice.flatten() - expected_slice).max() self.assertLessEqual(max_diff, 1e-3) diff --git a/tests/pipelines/pag/test_pag_sd_img2img.py b/tests/pipelines/pag/test_pag_sd_img2img.py index 0b440d5ec9fc..1e9b3c24c9ac 100644 --- a/tests/pipelines/pag/test_pag_sd_img2img.py +++ b/tests/pipelines/pag/test_pag_sd_img2img.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -206,7 +206,7 @@ def test_pag_inference(self): ), f"the shape of the output image should be (1, 32, 32, 3) but got {image.shape}" expected_slice = np.array( - [0.44203848, 0.49598145, 0.42248967, 0.6707724, 0.5683791, 0.43603387, 0.58316565, 0.60077155, 0.5174199] + [0.4508819, 0.49191576, 0.42664427, 0.66448534, 0.5606137, 0.43760118, 0.58251137, 0.5944448, 0.51642907] ) max_diff = np.abs(image_slice.flatten() - expected_slice).max() self.assertLessEqual(max_diff, 1e-3) diff --git a/tests/pipelines/pag/test_pag_sd_inpaint.py b/tests/pipelines/pag/test_pag_sd_inpaint.py index 754158bbf138..5a78ac6ade12 100644 --- a/tests/pipelines/pag/test_pag_sd_inpaint.py +++ b/tests/pipelines/pag/test_pag_sd_inpaint.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -239,7 +239,7 @@ def test_pag_inference(self): 3, ), f"the shape of the output image should be (1, 64, 64, 3) but got {image.shape}" - expected_slice = np.array([0.7190, 0.5807, 0.6007, 0.5600, 0.6350, 0.6639, 0.5680, 0.5664, 0.5230]) + expected_slice = np.array([0.7173, 0.5821, 0.6031, 0.5765, 0.6412, 0.6558, 0.5803, 0.5675, 0.5246]) max_diff = np.abs(image_slice.flatten() - expected_slice).max() assert max_diff < 1e-3, f"output is different from expected, {image_slice.flatten()}" diff --git a/tests/pipelines/pag/test_pag_sdxl.py b/tests/pipelines/pag/test_pag_sdxl.py index cca5c61651b3..af415ab71071 100644 --- a/tests/pipelines/pag/test_pag_sdxl.py +++ b/tests/pipelines/pag/test_pag_sdxl.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -279,7 +279,7 @@ def test_pag_inference(self): 64, 3, ), f"the shape of the output image should be (1, 64, 64, 3) but got {image.shape}" - expected_slice = np.array([0.5382, 0.5439, 0.4704, 0.4569, 0.5234, 0.4834, 0.5289, 0.5039, 0.4764]) + expected_slice = np.array([0.5565, 0.5305, 0.4652, 0.4330, 0.4823, 0.4640, 0.5191, 0.4983, 0.4684]) max_diff = np.abs(image_slice.flatten() - expected_slice).max() self.assertLessEqual(max_diff, 1e-3) diff --git a/tests/pipelines/pag/test_pag_sdxl_img2img.py b/tests/pipelines/pag/test_pag_sdxl_img2img.py index d311500d3ca7..821462d90d65 100644 --- a/tests/pipelines/pag/test_pag_sdxl_img2img.py +++ b/tests/pipelines/pag/test_pag_sdxl_img2img.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -258,7 +258,7 @@ def test_pag_inference(self): 32, 3, ), f"the shape of the output image should be (1, 64, 64, 3) but got {image.shape}" - expected_slice = np.array([0.4613, 0.4902, 0.4406, 0.6788, 0.5611, 0.4529, 0.5893, 0.5975, 0.5226]) + expected_slice = np.array([0.4566, 0.4907, 0.4374, 0.6633, 0.5626, 0.4494, 0.5771, 0.6011, 0.5245]) max_diff = np.abs(image_slice.flatten() - expected_slice).max() assert max_diff < 1e-3, f"output is different from expected, {image_slice.flatten()}" diff --git a/tests/pipelines/pag/test_pag_sdxl_inpaint.py b/tests/pipelines/pag/test_pag_sdxl_inpaint.py index 00a07582e205..0db295c44f71 100644 --- a/tests/pipelines/pag/test_pag_sdxl_inpaint.py +++ b/tests/pipelines/pag/test_pag_sdxl_inpaint.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -263,7 +263,7 @@ def test_pag_inference(self): 64, 3, ), f"the shape of the output image should be (1, 64, 64, 3) but got {image.shape}" - expected_slice = np.array([0.8366, 0.5513, 0.6105, 0.6213, 0.6957, 0.7400, 0.6614, 0.6102, 0.5239]) + expected_slice = np.array([0.7893, 0.5446, 0.5826, 0.6441, 0.6660, 0.7566, 0.6605, 0.5838, 0.5160]) max_diff = np.abs(image_slice.flatten() - expected_slice).max() assert max_diff < 1e-3, f"output is different from expected, {image_slice.flatten()}" diff --git a/tests/pipelines/pixart_alpha/test_pixart.py b/tests/pipelines/pixart_alpha/test_pixart.py index 399bcdf75609..b68841159b4d 100644 --- a/tests/pipelines/pixart_alpha/test_pixart.py +++ b/tests/pipelines/pixart_alpha/test_pixart.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/pipelines/pixart_sigma/test_pixart.py b/tests/pipelines/pixart_sigma/test_pixart.py index 51eebadd0ed0..a6f73ff50f00 100644 --- a/tests/pipelines/pixart_sigma/test_pixart.py +++ b/tests/pipelines/pixart_sigma/test_pixart.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/pipelines/pndm/test_pndm.py b/tests/pipelines/pndm/test_pndm.py index 61d6efe88ccd..5bca4bdede03 100644 --- a/tests/pipelines/pndm/test_pndm.py +++ b/tests/pipelines/pndm/test_pndm.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/pipelines/shap_e/test_shap_e.py b/tests/pipelines/shap_e/test_shap_e.py index 99fd28692981..4d08b72f7f32 100644 --- a/tests/pipelines/shap_e/test_shap_e.py +++ b/tests/pipelines/shap_e/test_shap_e.py @@ -1,4 +1,4 @@ -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -187,7 +187,7 @@ def test_shap_e(self): assert image.shape == (32, 16) - expected_slice = np.array([-1.0000, -0.6559, 1.0000, -0.9096, -0.7252, 0.8211, -0.7647, -0.3308, 0.6462]) + expected_slice = np.array([-1.0000, -0.6097, 1.0000, -0.9113, -0.6667, 0.8234, -0.7639, -0.2809, 0.6456]) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 def test_inference_batch_consistent(self): diff --git a/tests/pipelines/shap_e/test_shap_e_img2img.py b/tests/pipelines/shap_e/test_shap_e_img2img.py index b1867db249ea..406be60da0b4 100644 --- a/tests/pipelines/shap_e/test_shap_e_img2img.py +++ b/tests/pipelines/shap_e/test_shap_e_img2img.py @@ -1,4 +1,4 @@ -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -205,7 +205,7 @@ def test_shap_e(self): assert image.shape == (32, 16) expected_slice = np.array( - [-1.0, 0.40668195, 0.57322013, -0.9469888, 0.4283227, 0.30348337, -0.81094897, 0.74555075, 0.15342723] + [-1.0, -0.08255208, 1.0, -0.99469435, -0.19953543, 0.99999994, -0.9133135, 0.24063873, 0.8238962] ) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 diff --git a/tests/pipelines/stable_audio/test_stable_audio.py b/tests/pipelines/stable_audio/test_stable_audio.py index 492aa92252de..606de82324ca 100644 --- a/tests/pipelines/stable_audio/test_stable_audio.py +++ b/tests/pipelines/stable_audio/test_stable_audio.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/pipelines/stable_cascade/test_stable_cascade_combined.py b/tests/pipelines/stable_cascade/test_stable_cascade_combined.py index d9a511ab199c..9cb3fdf70c06 100644 --- a/tests/pipelines/stable_cascade/test_stable_cascade_combined.py +++ b/tests/pipelines/stable_cascade/test_stable_cascade_combined.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/pipelines/stable_cascade/test_stable_cascade_decoder.py b/tests/pipelines/stable_cascade/test_stable_cascade_decoder.py index b92df4c5d268..180b723de064 100644 --- a/tests/pipelines/stable_cascade/test_stable_cascade_decoder.py +++ b/tests/pipelines/stable_cascade/test_stable_cascade_decoder.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/pipelines/stable_cascade/test_stable_cascade_prior.py b/tests/pipelines/stable_cascade/test_stable_cascade_prior.py index 0bc821b7e64f..62bd84499ec9 100644 --- a/tests/pipelines/stable_cascade/test_stable_cascade_prior.py +++ b/tests/pipelines/stable_cascade/test_stable_cascade_prior.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/pipelines/stable_diffusion/test_onnx_stable_diffusion.py b/tests/pipelines/stable_diffusion/test_onnx_stable_diffusion.py index 62414f3f1947..77718dbf5ab9 100644 --- a/tests/pipelines/stable_diffusion/test_onnx_stable_diffusion.py +++ b/tests/pipelines/stable_diffusion/test_onnx_stable_diffusion.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/pipelines/stable_diffusion/test_onnx_stable_diffusion_img2img.py b/tests/pipelines/stable_diffusion/test_onnx_stable_diffusion_img2img.py index 28d1d0f37ff8..1a5ac9446285 100644 --- a/tests/pipelines/stable_diffusion/test_onnx_stable_diffusion_img2img.py +++ b/tests/pipelines/stable_diffusion/test_onnx_stable_diffusion_img2img.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/pipelines/stable_diffusion/test_onnx_stable_diffusion_inpaint.py b/tests/pipelines/stable_diffusion/test_onnx_stable_diffusion_inpaint.py index 1d46ff9a2f5f..80986ca313af 100644 --- a/tests/pipelines/stable_diffusion/test_onnx_stable_diffusion_inpaint.py +++ b/tests/pipelines/stable_diffusion/test_onnx_stable_diffusion_inpaint.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/pipelines/stable_diffusion/test_stable_diffusion.py b/tests/pipelines/stable_diffusion/test_stable_diffusion.py index c9d9525b2e45..d1fa29fb64ac 100644 --- a/tests/pipelines/stable_diffusion/test_stable_diffusion.py +++ b/tests/pipelines/stable_diffusion/test_stable_diffusion.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -181,7 +181,7 @@ def test_stable_diffusion_ddim(self): image_slice = image[0, -3:, -3:, -1] assert image.shape == (1, 64, 64, 3) - expected_slice = np.array([0.1763, 0.4776, 0.4986, 0.2566, 0.3802, 0.4596, 0.5363, 0.3277, 0.3949]) + expected_slice = np.array([0.1641, 0.4640, 0.4864, 0.2683, 0.3652, 0.4507, 0.5290, 0.3307, 0.3977]) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 @@ -356,7 +356,7 @@ def test_stable_diffusion_ddim_factor_8(self): image_slice = image[0, -3:, -3:, -1] assert image.shape == (1, 136, 136, 3) - expected_slice = np.array([0.4720, 0.5426, 0.5160, 0.3961, 0.4696, 0.4296, 0.5738, 0.5888, 0.5481]) + expected_slice = np.array([0.4587, 0.5238, 0.5006, 0.3869, 0.4538, 0.4228, 0.5666, 0.5814, 0.5468]) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 @@ -374,7 +374,7 @@ def test_stable_diffusion_pndm(self): image_slice = image[0, -3:, -3:, -1] assert image.shape == (1, 64, 64, 3) - expected_slice = np.array([0.1941, 0.4748, 0.4880, 0.2222, 0.4221, 0.4545, 0.5604, 0.3488, 0.3902]) + expected_slice = np.array([0.1791, 0.4611, 0.4741, 0.2336, 0.4175, 0.4484, 0.5582, 0.3556, 0.3951]) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 @@ -414,7 +414,7 @@ def test_stable_diffusion_k_lms(self): image_slice = image[0, -3:, -3:, -1] assert image.shape == (1, 64, 64, 3) - expected_slice = np.array([0.2681, 0.4785, 0.4857, 0.2426, 0.4473, 0.4481, 0.5610, 0.3676, 0.3855]) + expected_slice = np.array([0.2185, 0.4538, 0.4660, 0.2454, 0.4408, 0.4377, 0.5629, 0.3754, 0.3927]) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 @@ -433,7 +433,7 @@ def test_stable_diffusion_k_euler_ancestral(self): image_slice = image[0, -3:, -3:, -1] assert image.shape == (1, 64, 64, 3) - expected_slice = np.array([0.2682, 0.4782, 0.4855, 0.2424, 0.4472, 0.4479, 0.5612, 0.3676, 0.3854]) + expected_slice = np.array([0.2185, 0.4534, 0.4657, 0.2452, 0.4406, 0.4375, 0.5631, 0.3755, 0.3927]) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 @@ -452,7 +452,7 @@ def test_stable_diffusion_k_euler(self): image_slice = image[0, -3:, -3:, -1] assert image.shape == (1, 64, 64, 3) - expected_slice = np.array([0.2681, 0.4785, 0.4857, 0.2426, 0.4473, 0.4481, 0.5610, 0.3676, 0.3855]) + expected_slice = np.array([0.2185, 0.4538, 0.4660, 0.2454, 0.4408, 0.4377, 0.5629, 0.3754, 0.3927]) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 @@ -471,7 +471,7 @@ def test_stable_diffusion_vae_slicing(self): output_1 = sd_pipe(**inputs) # make sure sliced vae decode yields the same result - sd_pipe.enable_vae_slicing() + sd_pipe.vae.enable_slicing() inputs = self.get_dummy_inputs(device) inputs["prompt"] = [inputs["prompt"]] * image_count output_2 = sd_pipe(**inputs) @@ -496,7 +496,7 @@ def test_stable_diffusion_vae_tiling(self): output_1 = sd_pipe([prompt], generator=generator, guidance_scale=6.0, num_inference_steps=2, output_type="np") # make sure tiled vae decode yields the same result - sd_pipe.enable_vae_tiling() + sd_pipe.vae.enable_tiling() generator = torch.Generator(device=device).manual_seed(0) output_2 = sd_pipe([prompt], generator=generator, guidance_scale=6.0, num_inference_steps=2, output_type="np") @@ -524,7 +524,7 @@ def test_stable_diffusion_negative_prompt(self): image_slice = image[0, -3:, -3:, -1] assert image.shape == (1, 64, 64, 3) - expected_slice = np.array([0.1907, 0.4709, 0.4858, 0.2224, 0.4223, 0.4539, 0.5606, 0.3489, 0.3900]) + expected_slice = np.array([0.1772, 0.4578, 0.4700, 0.2363, 0.4173, 0.4462, 0.5595, 0.3588, 0.3959]) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 @@ -889,7 +889,7 @@ def test_stable_diffusion_vae_slicing(self): pipe.enable_attention_slicing() # enable vae slicing - pipe.enable_vae_slicing() + pipe.vae.enable_slicing() inputs = self.get_inputs(torch_device, dtype=torch.float16) inputs["prompt"] = [inputs["prompt"]] * 4 inputs["latents"] = torch.cat([inputs["latents"]] * 4) @@ -901,7 +901,7 @@ def test_stable_diffusion_vae_slicing(self): assert mem_bytes < 4e9 # disable vae slicing - pipe.disable_vae_slicing() + pipe.vae.disable_slicing() inputs = self.get_inputs(torch_device, dtype=torch.float16) inputs["prompt"] = [inputs["prompt"]] * 4 inputs["latents"] = torch.cat([inputs["latents"]] * 4) @@ -928,7 +928,7 @@ def test_stable_diffusion_vae_tiling(self): prompt = "a photograph of an astronaut riding a horse" # enable vae tiling - pipe.enable_vae_tiling() + pipe.vae.enable_tiling() pipe.enable_model_cpu_offload(device=torch_device) generator = torch.Generator(device="cpu").manual_seed(0) output_chunked = pipe( @@ -945,7 +945,7 @@ def test_stable_diffusion_vae_tiling(self): mem_bytes = backend_max_memory_allocated(torch_device) # disable vae tiling - pipe.disable_vae_tiling() + pipe.vae.disable_tiling() generator = torch.Generator(device="cpu").manual_seed(0) output = pipe( [prompt], diff --git a/tests/pipelines/stable_diffusion/test_stable_diffusion_img2img.py b/tests/pipelines/stable_diffusion/test_stable_diffusion_img2img.py index a0b7268b9dd4..fcaf21e1b04a 100644 --- a/tests/pipelines/stable_diffusion/test_stable_diffusion_img2img.py +++ b/tests/pipelines/stable_diffusion/test_stable_diffusion_img2img.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -163,7 +163,7 @@ def test_stable_diffusion_img2img_default_case(self): image_slice = image[0, -3:, -3:, -1] assert image.shape == (1, 32, 32, 3) - expected_slice = np.array([0.4555, 0.3216, 0.4049, 0.4620, 0.4618, 0.4126, 0.4122, 0.4629, 0.4579]) + expected_slice = np.array([0.4517, 0.3640, 0.4036, 0.4014, 0.4381, 0.3990, 0.3824, 0.4628, 0.4415]) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3 @@ -180,7 +180,7 @@ def test_stable_diffusion_img2img_default_case_lcm(self): image_slice = image[0, -3:, -3:, -1] assert image.shape == (1, 32, 32, 3) - expected_slice = np.array([0.5709, 0.4614, 0.4587, 0.5978, 0.5298, 0.6910, 0.6240, 0.5212, 0.5454]) + expected_slice = np.array([0.5915, 0.4625, 0.4547, 0.6083, 0.5414, 0.6927, 0.6253, 0.5223, 0.5449]) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3 @@ -199,7 +199,7 @@ def test_stable_diffusion_img2img_default_case_lcm_custom_timesteps(self): image_slice = image[0, -3:, -3:, -1] assert image.shape == (1, 32, 32, 3) - expected_slice = np.array([0.5709, 0.4614, 0.4587, 0.5978, 0.5298, 0.6910, 0.6240, 0.5212, 0.5454]) + expected_slice = np.array([0.5915, 0.4625, 0.4547, 0.6083, 0.5414, 0.6927, 0.6253, 0.5223, 0.5449]) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3 @@ -217,14 +217,14 @@ def test_stable_diffusion_img2img_negative_prompt(self): image_slice = image[0, -3:, -3:, -1] assert image.shape == (1, 32, 32, 3) - expected_slice = np.array([0.4593, 0.3408, 0.4232, 0.4749, 0.4476, 0.4115, 0.4357, 0.4733, 0.4663]) + expected_slice = np.array([0.4642, 0.3576, 0.4112, 0.4289, 0.4396, 0.4068, 0.3845, 0.4603, 0.4427]) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3 def test_ip_adapter(self): expected_pipe_slice = None if torch_device == "cpu": - expected_pipe_slice = np.array([0.4932, 0.5092, 0.5135, 0.5517, 0.5626, 0.6621, 0.6490, 0.5021, 0.5441]) + expected_pipe_slice = np.array([0.4923, 0.5012, 0.5041, 0.5462, 0.5665, 0.6621, 0.6451, 0.5046, 0.5431]) return super().test_ip_adapter(expected_pipe_slice=expected_pipe_slice) def test_stable_diffusion_img2img_multiple_init_images(self): @@ -241,7 +241,7 @@ def test_stable_diffusion_img2img_multiple_init_images(self): image_slice = image[-1, -3:, -3:, -1] assert image.shape == (2, 32, 32, 3) - expected_slice = np.array([0.4241, 0.5576, 0.5711, 0.4792, 0.4311, 0.5952, 0.5827, 0.5138, 0.5109]) + expected_slice = np.array([0.4385, 0.3589, 0.4905, 0.4607, 0.4040, 0.5551, 0.5114, 0.5269, 0.5267]) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3 @@ -260,7 +260,7 @@ def test_stable_diffusion_img2img_k_lms(self): image_slice = image[0, -3:, -3:, -1] assert image.shape == (1, 32, 32, 3) - expected_slice = np.array([0.4398, 0.4949, 0.4337, 0.6580, 0.5555, 0.4338, 0.5769, 0.5955, 0.5175]) + expected_slice = np.array([0.4409, 0.4952, 0.4357, 0.6574, 0.5583, 0.4375, 0.5777, 0.5960, 0.5168]) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3 @@ -277,7 +277,7 @@ def test_stable_diffusion_img2img_tiny_autoencoder(self): image_slice = image[0, -3:, -3:, -1] assert image.shape == (1, 32, 32, 3) - expected_slice = np.array([0.00669, 0.00669, 0.0, 0.00693, 0.00858, 0.0, 0.00567, 0.00515, 0.00125]) + expected_slice = np.array([0.00551, 0.01295, 0.01731, 0.00725, 0.01140, 0.01334, 0.00152, 0.00997, 0.01265]) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3 diff --git a/tests/pipelines/stable_diffusion/test_stable_diffusion_inpaint.py b/tests/pipelines/stable_diffusion/test_stable_diffusion_inpaint.py index 259806a9479c..a69f71c37530 100644 --- a/tests/pipelines/stable_diffusion/test_stable_diffusion_inpaint.py +++ b/tests/pipelines/stable_diffusion/test_stable_diffusion_inpaint.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -180,7 +180,7 @@ def test_stable_diffusion_inpaint(self): image_slice = image[0, -3:, -3:, -1] assert image.shape == (1, 64, 64, 3) - expected_slice = np.array([0.4703, 0.5697, 0.3879, 0.5470, 0.6042, 0.4413, 0.5078, 0.4728, 0.4469]) + expected_slice = np.array([0.4816, 0.5766, 0.3897, 0.5448, 0.5982, 0.4359, 0.5142, 0.4836, 0.4536]) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 @@ -197,7 +197,7 @@ def test_stable_diffusion_inpaint_lcm(self): image_slice = image[0, -3:, -3:, -1] assert image.shape == (1, 64, 64, 3) - expected_slice = np.array([0.4931, 0.5988, 0.4569, 0.5556, 0.6650, 0.5087, 0.5966, 0.5358, 0.5269]) + expected_slice = np.array([0.4979, 0.5907, 0.4608, 0.5595, 0.6673, 0.5135, 0.6035, 0.5464, 0.5325]) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 @@ -216,7 +216,7 @@ def test_stable_diffusion_inpaint_lcm_custom_timesteps(self): image_slice = image[0, -3:, -3:, -1] assert image.shape == (1, 64, 64, 3) - expected_slice = np.array([0.4931, 0.5988, 0.4569, 0.5556, 0.6650, 0.5087, 0.5966, 0.5358, 0.5269]) + expected_slice = np.array([0.4979, 0.5907, 0.4608, 0.5595, 0.6673, 0.5135, 0.6035, 0.5464, 0.5325]) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 @@ -352,15 +352,6 @@ def callback_on_step_end(pipe, i, t, callback_kwargs): # they should be the same assert torch.allclose(intermediate_latent, output_interrupted, atol=1e-4) - def test_ip_adapter(self, from_simple=False, expected_pipe_slice=None): - if not from_simple: - expected_pipe_slice = None - if torch_device == "cpu": - expected_pipe_slice = np.array( - [0.4390, 0.5452, 0.3772, 0.5448, 0.6031, 0.4480, 0.5194, 0.4687, 0.4640] - ) - return super().test_ip_adapter(expected_pipe_slice=expected_pipe_slice) - def test_encode_prompt_works_in_isolation(self): extra_required_param_value_dict = { "device": torch.device(torch_device).type, @@ -455,12 +446,6 @@ def get_dummy_inputs_2images(self, device, seed=0, img_res=64): } return inputs - def test_ip_adapter(self): - expected_pipe_slice = None - if torch_device == "cpu": - expected_pipe_slice = np.array([0.6345, 0.5395, 0.5611, 0.5403, 0.5830, 0.5855, 0.5193, 0.5443, 0.5211]) - return super().test_ip_adapter(from_simple=True, expected_pipe_slice=expected_pipe_slice) - def test_stable_diffusion_inpaint(self): device = "cpu" # ensure determinism for the device-dependent torch.Generator components = self.get_dummy_components() @@ -563,7 +548,7 @@ def test_stable_diffusion_inpaint_euler(self): assert image.shape == (1, 64, 64, 3) expected_slice = np.array( - [[0.6387283, 0.5564158, 0.58631873, 0.5539942, 0.5494673, 0.6461868, 0.5251618, 0.5497595, 0.5508756]] + [[0.6443541, 0.55819386, 0.58656645, 0.5574256, 0.55064464, 0.64331865, 0.5232371, 0.5466772, 0.5478098]] ) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-4 diff --git a/tests/pipelines/stable_diffusion/test_stable_diffusion_instruction_pix2pix.py b/tests/pipelines/stable_diffusion/test_stable_diffusion_instruction_pix2pix.py index 4758c5dab44b..0852965c2fe8 100644 --- a/tests/pipelines/stable_diffusion/test_stable_diffusion_instruction_pix2pix.py +++ b/tests/pipelines/stable_diffusion/test_stable_diffusion_instruction_pix2pix.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -150,7 +150,7 @@ def test_stable_diffusion_pix2pix_default_case(self): image = sd_pipe(**inputs).images image_slice = image[0, -3:, -3:, -1] assert image.shape == (1, 32, 32, 3) - expected_slice = np.array([0.7526, 0.3750, 0.4547, 0.6117, 0.5866, 0.5016, 0.4327, 0.5642, 0.4815]) + expected_slice = np.array([0.7024, 0.3769, 0.4392, 0.5858, 0.5620, 0.4827, 0.4174, 0.5506, 0.4806]) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3 @@ -168,7 +168,7 @@ def test_stable_diffusion_pix2pix_negative_prompt(self): image_slice = image[0, -3:, -3:, -1] assert image.shape == (1, 32, 32, 3) - expected_slice = np.array([0.7511, 0.3642, 0.4553, 0.6236, 0.5797, 0.5013, 0.4343, 0.5611, 0.4831]) + expected_slice = np.array([0.7059, 0.3731, 0.4428, 0.5922, 0.5649, 0.4871, 0.4203, 0.5507, 0.4809]) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3 @@ -192,7 +192,7 @@ def test_stable_diffusion_pix2pix_multiple_init_images(self): image_slice = image[-1, -3:, -3:, -1] assert image.shape == (2, 32, 32, 3) - expected_slice = np.array([0.5812, 0.5748, 0.5222, 0.5908, 0.5695, 0.7174, 0.6804, 0.5523, 0.5579]) + expected_slice = np.array([0.5710, 0.5834, 0.5128, 0.5892, 0.5773, 0.7047, 0.6797, 0.5411, 0.5536]) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3 @@ -211,7 +211,7 @@ def test_stable_diffusion_pix2pix_euler(self): image_slice = image[0, -3:, -3:, -1] assert image.shape == (1, 32, 32, 3) - expected_slice = np.array([0.7417, 0.3842, 0.4732, 0.5776, 0.5891, 0.5139, 0.4052, 0.5673, 0.4986]) + expected_slice = np.array([0.6674, 0.3879, 0.4447, 0.5375, 0.5464, 0.4881, 0.3896, 0.5467, 0.4923]) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3 diff --git a/tests/pipelines/stable_diffusion_2/test_stable_diffusion.py b/tests/pipelines/stable_diffusion_2/test_stable_diffusion.py index 3b2552b432d3..83ebbd8bb717 100644 --- a/tests/pipelines/stable_diffusion_2/test_stable_diffusion.py +++ b/tests/pipelines/stable_diffusion_2/test_stable_diffusion.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -169,7 +169,7 @@ def test_stable_diffusion_ddim(self): image_slice = image[0, -3:, -3:, -1] assert image.shape == (1, 64, 64, 3) - expected_slice = np.array([0.5753, 0.6113, 0.5005, 0.5036, 0.5464, 0.4725, 0.4982, 0.4865, 0.4861]) + expected_slice = np.array([0.5679, 0.6010, 0.5052, 0.5159, 0.5389, 0.4739, 0.5074, 0.4872, 0.4843]) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 @@ -203,7 +203,7 @@ def test_stable_diffusion_k_lms(self): image_slice = image[0, -3:, -3:, -1] assert image.shape == (1, 64, 64, 3) - expected_slice = np.array([0.4865, 0.5439, 0.4840, 0.4995, 0.5543, 0.4846, 0.5199, 0.4942, 0.5061]) + expected_slice = np.array([0.4767, 0.5389, 0.4766, 0.5015, 0.5585, 0.4859, 0.5341, 0.5041, 0.5073]) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 @@ -220,7 +220,7 @@ def test_stable_diffusion_k_euler_ancestral(self): image_slice = image[0, -3:, -3:, -1] assert image.shape == (1, 64, 64, 3) - expected_slice = np.array([0.4864, 0.5440, 0.4842, 0.4994, 0.5543, 0.4846, 0.5196, 0.4942, 0.5063]) + expected_slice = np.array([0.4766, 0.5390, 0.4768, 0.5013, 0.5585, 0.4859, 0.5338, 0.5041, 0.5075]) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 @@ -237,7 +237,7 @@ def test_stable_diffusion_k_euler(self): image_slice = image[0, -3:, -3:, -1] assert image.shape == (1, 64, 64, 3) - expected_slice = np.array([0.4865, 0.5439, 0.4840, 0.4995, 0.5543, 0.4846, 0.5199, 0.4942, 0.5061]) + expected_slice = np.array([0.4767, 0.5389, 0.4766, 0.5015, 0.5585, 0.4859, 0.5341, 0.5041, 0.5073]) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 diff --git a/tests/pipelines/stable_diffusion_2/test_stable_diffusion_depth.py b/tests/pipelines/stable_diffusion_2/test_stable_diffusion_depth.py index bea7c099046f..dbd62836cc05 100644 --- a/tests/pipelines/stable_diffusion_2/test_stable_diffusion_depth.py +++ b/tests/pipelines/stable_diffusion_2/test_stable_diffusion_depth.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -297,7 +297,7 @@ def test_stable_diffusion_depth2img_default_case(self): if torch_device == "mps": expected_slice = np.array([0.6071, 0.5035, 0.4378, 0.5776, 0.5753, 0.4316, 0.4513, 0.5263, 0.4546]) else: - expected_slice = np.array([0.5435, 0.4992, 0.3783, 0.4411, 0.5842, 0.4654, 0.3786, 0.5077, 0.4655]) + expected_slice = np.array([0.4082, 0.5014, 0.4256, 0.3948, 0.4573, 0.4061, 0.3515, 0.4656, 0.4529]) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3 @@ -318,7 +318,7 @@ def test_stable_diffusion_depth2img_negative_prompt(self): if torch_device == "mps": expected_slice = np.array([0.6296, 0.5125, 0.3890, 0.4456, 0.5955, 0.4621, 0.3810, 0.5310, 0.4626]) else: - expected_slice = np.array([0.6012, 0.4507, 0.3769, 0.4121, 0.5566, 0.4585, 0.3803, 0.5045, 0.4631]) + expected_slice = np.array([0.4168, 0.4955, 0.4142, 0.4035, 0.4646, 0.4082, 0.3500, 0.4635, 0.4475]) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3 @@ -340,7 +340,7 @@ def test_stable_diffusion_depth2img_multiple_init_images(self): if torch_device == "mps": expected_slice = np.array([0.6501, 0.5150, 0.4939, 0.6688, 0.5437, 0.5758, 0.5115, 0.4406, 0.4551]) else: - expected_slice = np.array([0.6557, 0.6214, 0.6254, 0.5775, 0.4785, 0.5949, 0.5904, 0.4785, 0.4730]) + expected_slice = np.array([0.6261, 0.5452, 0.5330, 0.5160, 0.4768, 0.5216, 0.4413, 0.4797, 0.4777]) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3 @@ -359,7 +359,7 @@ def test_stable_diffusion_depth2img_pil(self): if torch_device == "mps": expected_slice = np.array([0.53232, 0.47015, 0.40868, 0.45651, 0.4891, 0.4668, 0.4287, 0.48822, 0.47439]) else: - expected_slice = np.array([0.5435, 0.4992, 0.3783, 0.4411, 0.5842, 0.4654, 0.3786, 0.5077, 0.4655]) + expected_slice = np.array([0.4082, 0.5014, 0.4256, 0.3948, 0.4573, 0.4061, 0.3515, 0.4656, 0.4529]) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3 diff --git a/tests/pipelines/stable_diffusion_2/test_stable_diffusion_inpaint.py b/tests/pipelines/stable_diffusion_2/test_stable_diffusion_inpaint.py index f010c1b03fe3..c1b8e65de787 100644 --- a/tests/pipelines/stable_diffusion_2/test_stable_diffusion_inpaint.py +++ b/tests/pipelines/stable_diffusion_2/test_stable_diffusion_inpaint.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -149,7 +149,7 @@ def test_stable_diffusion_inpaint(self): image_slice = image[0, -3:, -3:, -1] assert image.shape == (1, 64, 64, 3) - expected_slice = np.array([0.4727, 0.5735, 0.3941, 0.5446, 0.5926, 0.4394, 0.5062, 0.4654, 0.4476]) + expected_slice = np.array([0.4858, 0.5739, 0.3934, 0.5437, 0.5925, 0.4369, 0.5121, 0.4742, 0.4538]) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 diff --git a/tests/pipelines/stable_diffusion_2/test_stable_diffusion_latent_upscale.py b/tests/pipelines/stable_diffusion_2/test_stable_diffusion_latent_upscale.py index 285c2fea7ebc..9be555935d84 100644 --- a/tests/pipelines/stable_diffusion_2/test_stable_diffusion_latent_upscale.py +++ b/tests/pipelines/stable_diffusion_2/test_stable_diffusion_latent_upscale.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -175,7 +175,7 @@ def test_inference(self): self.assertEqual(image.shape, (1, 256, 256, 3)) expected_slice = np.array( - [0.47222412, 0.41921633, 0.44717434, 0.46874192, 0.42588258, 0.46150726, 0.4677534, 0.45583832, 0.48579055] + [0.57231164, 0.54510796, 0.5154225, 0.62793314, 0.58053195, 0.5522338, 0.6212865, 0.55557936, 0.5478562] ) max_diff = np.abs(image_slice.flatten() - expected_slice).max() self.assertLessEqual(max_diff, 1e-3) @@ -195,7 +195,7 @@ def test_stable_diffusion_latent_upscaler_negative_prompt(self): assert image.shape == (1, 256, 256, 3) expected_slice = np.array( - [0.43865365, 0.404124, 0.42618454, 0.44333526, 0.40564927, 0.43818694, 0.4411913, 0.43404633, 0.46392226] + [0.5130085, 0.5026671, 0.53117144, 0.6296477, 0.5712023, 0.6032776, 0.68555176, 0.58478534, 0.5989826] ) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3 @@ -215,7 +215,7 @@ def test_stable_diffusion_latent_upscaler_multiple_init_images(self): assert image.shape == (2, 256, 256, 3) expected_slice = np.array( - [0.38730142, 0.35695046, 0.40646142, 0.40967226, 0.3981609, 0.4195988, 0.4248805, 0.430259, 0.45694894] + [0.49870333, 0.49428767, 0.48577496, 0.5712294, 0.55280155, 0.5219455, 0.62869453, 0.5554026, 0.52685684] ) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3 diff --git a/tests/pipelines/stable_diffusion_2/test_stable_diffusion_upscale.py b/tests/pipelines/stable_diffusion_2/test_stable_diffusion_upscale.py index 481ac7f2d10f..a186710487d7 100644 --- a/tests/pipelines/stable_diffusion_2/test_stable_diffusion_upscale.py +++ b/tests/pipelines/stable_diffusion_2/test_stable_diffusion_upscale.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -174,7 +174,7 @@ def test_stable_diffusion_upscale(self): expected_height_width = low_res_image.size[0] * 4 assert image.shape == (1, expected_height_width, expected_height_width, 3) - expected_slice = np.array([0.3113, 0.3910, 0.4272, 0.4859, 0.5061, 0.4652, 0.5362, 0.5715, 0.5661]) + expected_slice = np.array([0.2631, 0.4038, 0.4338, 0.4254, 0.5002, 0.4831, 0.5073, 0.5619, 0.5597]) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 assert np.abs(image_from_tuple_slice.flatten() - expected_slice).max() < 1e-2 @@ -290,7 +290,7 @@ def test_stable_diffusion_upscale_prompt_embeds(self): expected_height_width = low_res_image.size[0] * 4 assert image.shape == (1, expected_height_width, expected_height_width, 3) - expected_slice = np.array([0.3113, 0.3910, 0.4272, 0.4859, 0.5061, 0.4652, 0.5362, 0.5715, 0.5661]) + expected_slice = np.array([0.2631, 0.4038, 0.4338, 0.4254, 0.5002, 0.4831, 0.5073, 0.5619, 0.5597]) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 assert np.abs(image_from_prompt_embeds_slice.flatten() - expected_slice).max() < 1e-2 diff --git a/tests/pipelines/stable_diffusion_2/test_stable_diffusion_v_pred.py b/tests/pipelines/stable_diffusion_2/test_stable_diffusion_v_pred.py index 37b309c4cac4..f2011abdb0e9 100644 --- a/tests/pipelines/stable_diffusion_2/test_stable_diffusion_v_pred.py +++ b/tests/pipelines/stable_diffusion_2/test_stable_diffusion_v_pred.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -164,7 +164,7 @@ def test_stable_diffusion_v_pred_ddim(self): image_from_tuple_slice = image_from_tuple[0, -3:, -3:, -1] assert image.shape == (1, 64, 64, 3) - expected_slice = np.array([0.6569, 0.6525, 0.5142, 0.4968, 0.4923, 0.4601, 0.4996, 0.5041, 0.4544]) + expected_slice = np.array([0.6308, 0.6095, 0.5343, 0.5073, 0.4634, 0.4632, 0.4967, 0.4885, 0.4489]) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 assert np.abs(image_from_tuple_slice.flatten() - expected_slice).max() < 1e-2 @@ -214,7 +214,7 @@ def test_stable_diffusion_v_pred_k_euler(self): image_from_tuple_slice = image_from_tuple[0, -3:, -3:, -1] assert image.shape == (1, 64, 64, 3) - expected_slice = np.array([0.5644, 0.6514, 0.5190, 0.5663, 0.5287, 0.4953, 0.5430, 0.5243, 0.4778]) + expected_slice = np.array([0.5419, 0.5719, 0.5558, 0.5034, 0.4559, 0.5026, 0.4736, 0.4782, 0.4725]) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 assert np.abs(image_from_tuple_slice.flatten() - expected_slice).max() < 1e-2 diff --git a/tests/pipelines/stable_diffusion_3/test_pipeline_stable_diffusion_3.py b/tests/pipelines/stable_diffusion_3/test_pipeline_stable_diffusion_3.py index 200c832d0941..039b7924d017 100644 --- a/tests/pipelines/stable_diffusion_3/test_pipeline_stable_diffusion_3.py +++ b/tests/pipelines/stable_diffusion_3/test_pipeline_stable_diffusion_3.py @@ -143,7 +143,7 @@ def test_inference(self): generated_slice = np.concatenate([generated_slice[:8], generated_slice[-8:]]) # fmt: off - expected_slice = np.array([0.5112, 0.5228, 0.5235, 0.5524, 0.3188, 0.5017, 0.5574, 0.4899, 0.6812, 0.5991, 0.3908, 0.5213, 0.5582, 0.4457, 0.4204, 0.5616]) + expected_slice = np.array([0.5034, 0.4693, 0.4877, 0.5052, 0.2697, 0.4615, 0.5329, 0.4330, 0.6825, 0.5969, 0.3924, 0.5542, 0.5454, 0.4295, 0.4178, 0.4805]) # fmt: on self.assertTrue( diff --git a/tests/pipelines/stable_diffusion_3/test_pipeline_stable_diffusion_3_img2img.py b/tests/pipelines/stable_diffusion_3/test_pipeline_stable_diffusion_3_img2img.py index 3f46b341a09e..ef31cbfe0108 100644 --- a/tests/pipelines/stable_diffusion_3/test_pipeline_stable_diffusion_3_img2img.py +++ b/tests/pipelines/stable_diffusion_3/test_pipeline_stable_diffusion_3_img2img.py @@ -147,7 +147,7 @@ def test_inference(self): generated_slice = np.concatenate([generated_slice[:8], generated_slice[-8:]]) # fmt: off - expected_slice = np.array([0.4564, 0.5486, 0.4868, 0.5923, 0.3775, 0.5543, 0.4807, 0.4177, 0.3778, 0.5957, 0.5726, 0.4333, 0.6312, 0.5062, 0.4838, 0.5984]) + expected_slice = np.array([0.4301, 0.4957, 0.4331, 0.5094, 0.3746, 0.4505, 0.4512, 0.3695, 0.3894, 0.5979, 0.5911, 0.4989, 0.5575, 0.5186, 0.5015, 0.5804]) # fmt: on self.assertTrue( diff --git a/tests/pipelines/stable_diffusion_3/test_pipeline_stable_diffusion_3_inpaint.py b/tests/pipelines/stable_diffusion_3/test_pipeline_stable_diffusion_3_inpaint.py index a90ca21a801b..6c0d2b1a832e 100644 --- a/tests/pipelines/stable_diffusion_3/test_pipeline_stable_diffusion_3_inpaint.py +++ b/tests/pipelines/stable_diffusion_3/test_pipeline_stable_diffusion_3_inpaint.py @@ -151,7 +151,7 @@ def test_inference(self): generated_slice = np.concatenate([generated_slice[:8], generated_slice[-8:]]) # fmt: off - expected_slice = np.array([0.5035, 0.6661, 0.5859, 0.413, 0.4224, 0.4234, 0.7181, 0.5062, 0.5183, 0.6877, 0.5074, 0.585, 0.6111, 0.5422, 0.5306, 0.5891]) + expected_slice = np.array([0.5076, 0.6669, 0.5967, 0.4238, 0.3966, 0.4218, 0.7243, 0.4943, 0.5130, 0.7069, 0.5334, 0.5743, 0.6149, 0.5476, 0.5295, 0.6009]) # fmt: on self.assertTrue( diff --git a/tests/pipelines/stable_diffusion_adapter/test_stable_diffusion_adapter.py b/tests/pipelines/stable_diffusion_adapter/test_stable_diffusion_adapter.py index 79b38d1cad1c..d436758a965d 100644 --- a/tests/pipelines/stable_diffusion_adapter/test_stable_diffusion_adapter.py +++ b/tests/pipelines/stable_diffusion_adapter/test_stable_diffusion_adapter.py @@ -42,6 +42,8 @@ load_image, load_numpy, numpy_cosine_similarity_distance, + require_accelerate_version_greater, + require_accelerator, require_torch_accelerator, slow, torch_device, @@ -311,7 +313,7 @@ def test_adapter_lcm(self): image_slice = image[0, -3:, -3:, -1] assert image.shape == (1, 64, 64, 3) - expected_slice = np.array([0.4535, 0.5493, 0.4359, 0.5452, 0.6086, 0.4441, 0.5544, 0.501, 0.4859]) + expected_slice = np.array([0.4532, 0.5410, 0.4295, 0.5327, 0.6015, 0.4396, 0.5432, 0.4957, 0.4827]) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 @@ -333,7 +335,7 @@ def test_adapter_lcm_custom_timesteps(self): image_slice = image[0, -3:, -3:, -1] assert image.shape == (1, 64, 64, 3) - expected_slice = np.array([0.4535, 0.5493, 0.4359, 0.5452, 0.6086, 0.4441, 0.5544, 0.501, 0.4859]) + expected_slice = np.array([0.4532, 0.5410, 0.4295, 0.5327, 0.6015, 0.4396, 0.5432, 0.4957, 0.4827]) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 @@ -366,9 +368,11 @@ def test_stable_diffusion_adapter_default_case(self): image_slice = image[0, -3:, -3:, -1] assert image.shape == (1, 64, 64, 3) - expected_slice = np.array([0.4858, 0.5500, 0.4278, 0.4669, 0.6184, 0.4322, 0.5010, 0.5033, 0.4746]) + expected_slice = np.array([0.5248, 0.5794, 0.4504, 0.4649, 0.6327, 0.4491, 0.4922, 0.5155, 0.4938]) assert np.abs(image_slice.flatten() - expected_slice).max() < 5e-3 + @require_accelerator + @require_accelerate_version_greater("0.14.0") def test_from_pipe_consistent_forward_pass_cpu_offload(self): super().test_from_pipe_consistent_forward_pass_cpu_offload(expected_max_diff=6e-3) @@ -392,7 +396,7 @@ def test_stable_diffusion_adapter_default_case(self): image_slice = image[0, -3:, -3:, -1] assert image.shape == (1, 64, 64, 3) - expected_slice = np.array([0.4965, 0.5548, 0.4330, 0.4771, 0.6226, 0.4382, 0.5037, 0.5071, 0.4782]) + expected_slice = np.array([0.5463, 0.5897, 0.4547, 0.4751, 0.6357, 0.4527, 0.4924, 0.5190, 0.4969]) assert np.abs(image_slice.flatten() - expected_slice).max() < 5e-3 @@ -422,7 +426,7 @@ def test_stable_diffusion_adapter_default_case(self): image_slice = image[0, -3:, -3:, -1] assert image.shape == (1, 64, 64, 3) - expected_slice = np.array([0.4902, 0.5539, 0.4317, 0.4682, 0.6190, 0.4351, 0.5018, 0.5046, 0.4772]) + expected_slice = np.array([0.5368, 0.5864, 0.4573, 0.4682, 0.6317, 0.4550, 0.4931, 0.5175, 0.4986]) assert np.abs(image_slice.flatten() - expected_slice).max() < 5e-3 def test_inference_batch_consistent( diff --git a/tests/pipelines/stable_diffusion_image_variation/test_stable_diffusion_image_variation.py b/tests/pipelines/stable_diffusion_image_variation/test_stable_diffusion_image_variation.py index dbf5a7b68eae..7bc9269d02f2 100644 --- a/tests/pipelines/stable_diffusion_image_variation/test_stable_diffusion_image_variation.py +++ b/tests/pipelines/stable_diffusion_image_variation/test_stable_diffusion_image_variation.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -139,7 +139,7 @@ def test_stable_diffusion_img_variation_default_case(self): image_slice = image[0, -3:, -3:, -1] assert image.shape == (1, 64, 64, 3) - expected_slice = np.array([0.5348, 0.5924, 0.4798, 0.5237, 0.5741, 0.4651, 0.5344, 0.4942, 0.4851]) + expected_slice = np.array([0.5157, 0.5653, 0.4707, 0.4792, 0.5563, 0.4569, 0.5020, 0.4858, 0.4940]) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3 @@ -159,7 +159,7 @@ def test_stable_diffusion_img_variation_multiple_images(self): image_slice = image[-1, -3:, -3:, -1] assert image.shape == (2, 64, 64, 3) - expected_slice = np.array([0.6647, 0.5557, 0.5723, 0.5567, 0.5869, 0.6044, 0.5502, 0.5439, 0.5189]) + expected_slice = np.array([0.6749, 0.5424, 0.5692, 0.5694, 0.6228, 0.6268, 0.5516, 0.5598, 0.5227]) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3 diff --git a/tests/pipelines/stable_diffusion_xl/test_stable_diffusion_xl.py b/tests/pipelines/stable_diffusion_xl/test_stable_diffusion_xl.py index 7cf1f7dec3bc..2115b8a6f3f2 100644 --- a/tests/pipelines/stable_diffusion_xl/test_stable_diffusion_xl.py +++ b/tests/pipelines/stable_diffusion_xl/test_stable_diffusion_xl.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -179,7 +179,7 @@ def test_stable_diffusion_xl_euler(self): image_slice = image[0, -3:, -3:, -1] assert image.shape == (1, 64, 64, 3) - expected_slice = np.array([0.5388, 0.5452, 0.4694, 0.4583, 0.5253, 0.4832, 0.5288, 0.5035, 0.47]) + expected_slice = np.array([0.5556, 0.5318, 0.4649, 0.4321, 0.4824, 0.4624, 0.5183, 0.4981, 0.4692]) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 @@ -260,7 +260,7 @@ def test_stable_diffusion_ays(self): def test_ip_adapter(self): expected_pipe_slice = None if torch_device == "cpu": - expected_pipe_slice = np.array([0.5388, 0.5452, 0.4694, 0.4583, 0.5253, 0.4832, 0.5288, 0.5035, 0.4766]) + expected_pipe_slice = np.array([0.5556, 0.5318, 0.4649, 0.4321, 0.4824, 0.4624, 0.5183, 0.4981, 0.4692]) return super().test_ip_adapter(expected_pipe_slice=expected_pipe_slice) diff --git a/tests/pipelines/stable_diffusion_xl/test_stable_diffusion_xl_adapter.py b/tests/pipelines/stable_diffusion_xl/test_stable_diffusion_xl_adapter.py index 3d72270dda5c..132cd9c85842 100644 --- a/tests/pipelines/stable_diffusion_xl/test_stable_diffusion_xl_adapter.py +++ b/tests/pipelines/stable_diffusion_xl/test_stable_diffusion_xl_adapter.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -292,7 +292,7 @@ def test_ip_adapter(self, from_multi=False, expected_pipe_slice=None): expected_pipe_slice = None if torch_device == "cpu": expected_pipe_slice = np.array( - [0.5752, 0.6155, 0.4826, 0.5111, 0.5741, 0.4678, 0.5199, 0.5231, 0.4794] + [0.6002, 0.6262, 0.4981, 0.5304, 0.5774, 0.4685, 0.5228, 0.5208, 0.4938] ) return super().test_ip_adapter(expected_pipe_slice=expected_pipe_slice) @@ -313,7 +313,7 @@ def test_stable_diffusion_adapter_default_case(self): image_slice = image[0, -3:, -3:, -1] assert image.shape == (1, 64, 64, 3) - expected_slice = np.array([00.5752, 0.6155, 0.4826, 0.5111, 0.5741, 0.4678, 0.5199, 0.5231, 0.4794]) + expected_slice = np.array([0.6002, 0.6262, 0.4981, 0.5304, 0.5774, 0.4685, 0.5228, 0.5208, 0.4938]) assert np.abs(image_slice.flatten() - expected_slice).max() < 5e-3 @parameterized.expand( @@ -445,13 +445,13 @@ def test_stable_diffusion_adapter_default_case(self): image_slice = image[0, -3:, -3:, -1] assert image.shape == (1, 64, 64, 3) - expected_slice = np.array([0.5617, 0.6081, 0.4807, 0.5071, 0.5665, 0.4614, 0.5165, 0.5164, 0.4786]) + expected_slice = np.array([0.6114, 0.6256, 0.4972, 0.5219, 0.5668, 0.4658, 0.5210, 0.5188, 0.4908]) assert np.abs(image_slice.flatten() - expected_slice).max() < 5e-3 def test_ip_adapter(self): expected_pipe_slice = None if torch_device == "cpu": - expected_pipe_slice = np.array([0.5617, 0.6081, 0.4807, 0.5071, 0.5665, 0.4614, 0.5165, 0.5164, 0.4786]) + expected_pipe_slice = np.array([0.6114, 0.6256, 0.4972, 0.5219, 0.5668, 0.4658, 0.5210, 0.5188, 0.4908]) return super().test_ip_adapter(from_multi=True, expected_pipe_slice=expected_pipe_slice) diff --git a/tests/pipelines/stable_diffusion_xl/test_stable_diffusion_xl_img2img.py b/tests/pipelines/stable_diffusion_xl/test_stable_diffusion_xl_img2img.py index c5499847069f..9d2a6f0714c2 100644 --- a/tests/pipelines/stable_diffusion_xl/test_stable_diffusion_xl_img2img.py +++ b/tests/pipelines/stable_diffusion_xl/test_stable_diffusion_xl_img2img.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -218,7 +218,7 @@ def test_stable_diffusion_xl_img2img_euler(self): assert image.shape == (1, 32, 32, 3) - expected_slice = np.array([0.4664, 0.4886, 0.4403, 0.6902, 0.5592, 0.4534, 0.5931, 0.5951, 0.5224]) + expected_slice = np.array([0.4555, 0.4937, 0.4329, 0.6746, 0.5620, 0.4489, 0.5854, 0.5960, 0.5203]) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 @@ -236,7 +236,7 @@ def test_stable_diffusion_xl_img2img_euler_lcm(self): assert image.shape == (1, 32, 32, 3) - expected_slice = np.array([0.5604, 0.4352, 0.4717, 0.5844, 0.5101, 0.6704, 0.6290, 0.5460, 0.5286]) + expected_slice = np.array([0.5659, 0.4335, 0.4640, 0.5914, 0.5211, 0.6672, 0.6317, 0.5499, 0.5259]) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 @@ -256,7 +256,7 @@ def test_stable_diffusion_xl_img2img_euler_lcm_custom_timesteps(self): assert image.shape == (1, 32, 32, 3) - expected_slice = np.array([0.5604, 0.4352, 0.4717, 0.5844, 0.5101, 0.6704, 0.6290, 0.5460, 0.5286]) + expected_slice = np.array([0.5659, 0.4335, 0.4640, 0.5914, 0.5211, 0.6672, 0.6317, 0.5499, 0.5259]) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 @@ -273,7 +273,7 @@ def test_save_load_optional_components(self): def test_ip_adapter(self): expected_pipe_slice = None if torch_device == "cpu": - expected_pipe_slice = np.array([0.5133, 0.4626, 0.4970, 0.6273, 0.5160, 0.6891, 0.6639, 0.5892, 0.5709]) + expected_pipe_slice = np.array([0.4783, 0.4636, 0.4970, 0.5885, 0.5135, 0.6950, 0.6465, 0.5758, 0.5723]) return super().test_ip_adapter(expected_pipe_slice=expected_pipe_slice) @@ -290,7 +290,7 @@ def test_stable_diffusion_xl_img2img_tiny_autoencoder(self): image_slice = image[0, -3:, -3:, -1].flatten() assert image.shape == (1, 32, 32, 3) - expected_slice = np.array([0.0, 0.0, 0.0106, 0.0, 0.0, 0.0087, 0.0052, 0.0062, 0.0177]) + expected_slice = np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]) assert np.allclose(image_slice, expected_slice, atol=1e-4, rtol=1e-4) @@ -593,7 +593,7 @@ def test_stable_diffusion_xl_img2img_euler(self): assert image.shape == (1, 32, 32, 3) - expected_slice = np.array([0.4745, 0.4924, 0.4338, 0.6468, 0.5547, 0.4419, 0.5646, 0.5897, 0.5146]) + expected_slice = np.array([0.4692, 0.4957, 0.4337, 0.6666, 0.5656, 0.4461, 0.5742, 0.5932, 0.5162]) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 diff --git a/tests/pipelines/stable_diffusion_xl/test_stable_diffusion_xl_inpaint.py b/tests/pipelines/stable_diffusion_xl/test_stable_diffusion_xl_inpaint.py index d3f5779c7633..2cda0f52d749 100644 --- a/tests/pipelines/stable_diffusion_xl/test_stable_diffusion_xl_inpaint.py +++ b/tests/pipelines/stable_diffusion_xl/test_stable_diffusion_xl_inpaint.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -234,7 +234,7 @@ def get_dummy_inputs_2images(self, device, seed=0, img_res=64): def test_ip_adapter(self): expected_pipe_slice = None if torch_device == "cpu": - expected_pipe_slice = np.array([0.8274, 0.5538, 0.6141, 0.5843, 0.6865, 0.7082, 0.5861, 0.6123, 0.5344]) + expected_pipe_slice = np.array([0.7629, 0.5586, 0.5911, 0.5941, 0.6568, 0.7128, 0.5961, 0.5896, 0.5277]) return super().test_ip_adapter(expected_pipe_slice=expected_pipe_slice) @@ -259,7 +259,7 @@ def test_stable_diffusion_xl_inpaint_euler(self): assert image.shape == (1, 64, 64, 3) - expected_slice = np.array([0.8279, 0.5673, 0.6088, 0.6156, 0.6923, 0.7347, 0.6547, 0.6108, 0.5198]) + expected_slice = np.array([0.7819, 0.5553, 0.5853, 0.6391, 0.6660, 0.7527, 0.6580, 0.5877, 0.5146]) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 @@ -395,7 +395,7 @@ def test_stable_diffusion_xl_refiner(self): assert image.shape == (1, 64, 64, 3) - expected_slice = np.array([0.7540, 0.5231, 0.5833, 0.6217, 0.6339, 0.7067, 0.6507, 0.5672, 0.5030]) + expected_slice = np.array([0.7432, 0.5123, 0.5937, 0.6343, 0.6309, 0.7179, 0.6593, 0.5638, 0.5031]) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 diff --git a/tests/pipelines/stable_unclip/test_stable_unclip_img2img.py b/tests/pipelines/stable_unclip/test_stable_unclip_img2img.py index e7a0fbccef67..261c1785b45e 100644 --- a/tests/pipelines/stable_unclip/test_stable_unclip_img2img.py +++ b/tests/pipelines/stable_unclip/test_stable_unclip_img2img.py @@ -188,7 +188,7 @@ def test_image_embeds_none(self): image_slice = image[0, -3:, -3:, -1] assert image.shape == (1, 32, 32, 3) - expected_slice = np.array([0.4397, 0.7080, 0.5590, 0.4255, 0.7181, 0.5938, 0.4051, 0.3720, 0.5116]) + expected_slice = np.array([0.4661, 0.8109, 0.6731, 0.4309, 0.6674, 0.7538, 0.3841, 0.4951, 0.5243]) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3 diff --git a/tests/pipelines/test_pipeline_utils.py b/tests/pipelines/test_pipeline_utils.py index 6d9e68197976..e0c73b96b8f9 100644 --- a/tests/pipelines/test_pipeline_utils.py +++ b/tests/pipelines/test_pipeline_utils.py @@ -218,6 +218,44 @@ def test_diffusers_is_compatible_no_components_only_variants(self): ] self.assertFalse(is_safetensors_compatible(filenames)) + def test_diffusers_is_compatible_no_components_safetensors(self): + filenames = [ + "diffusion_pytorch_model.safetensors", + ] + self.assertTrue(is_safetensors_compatible(filenames)) + + def test_diffusers_is_compatible_no_components_safetensors_only_variants(self): + filenames = [ + "diffusion_pytorch_model.fp16.safetensors", + ] + self.assertTrue(is_safetensors_compatible(filenames, variant="fp16")) + + def test_transformers_is_compatible_weightless_subfolder(self): + # transformers-style flat layout: transformers-named weights at the root + a weight-less subfolder + filenames = [ + "model.safetensors", + "scheduler/scheduler_config.json", + ] + self.assertTrue(is_safetensors_compatible(filenames)) + + def test_transformers_is_not_compatible_weightless_subfolder(self): + # same flat layout but only .bin weights at the root -> not safetensors compatible + filenames = [ + "pytorch_model.bin", + "scheduler/scheduler_config.json", + ] + self.assertFalse(is_safetensors_compatible(filenames)) + + def test_transformers_is_compatible_sharded_root_weights(self): + # sharded transformers-style weights at the repo root (e.g. DiffusionGemma's model-00001-of-00011.safetensors) + filenames = [ + "model-00001-of-00002.safetensors", + "model-00002-of-00002.safetensors", + "model.safetensors.index.json", + "scheduler/scheduler_config.json", + ] + self.assertTrue(is_safetensors_compatible(filenames)) + def test_is_compatible_mixed_variants(self): filenames = [ "unet/diffusion_pytorch_model.fp16.safetensors", diff --git a/tests/pipelines/test_pipelines.py b/tests/pipelines/test_pipelines.py index 1df2cfa569e7..a193fb85182d 100644 --- a/tests/pipelines/test_pipelines.py +++ b/tests/pipelines/test_pipelines.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -2083,6 +2083,28 @@ def test_smart_download(self): # is not downloaded, but all the expected ones assert not os.path.isfile(os.path.join(snapshot_dir, "big_array.npy")) + def test_download_flat_transformers_style_repo(self): + # Repos with a flat, transformers-style layout host a component's files at the repo root instead of in a + # subfolder (here `model` and `processor`; only `scheduler/` has a folder). The download patterns must + # pick up the transformers auxiliary files at the root, while unrelated root files are still skipped. + model_id = "hf-internal-testing/tiny-flat-transformers-style-pipe" + with tempfile.TemporaryDirectory() as tmpdirname: + snapshot_dir = DiffusionPipeline.download(model_id, cache_dir=tmpdirname, force_download=True) + + assert os.path.isfile(os.path.join(snapshot_dir, "model.safetensors")) + assert os.path.isfile(os.path.join(snapshot_dir, CONFIG_NAME)) + for aux_file in [ + "tokenizer.json", + "tokenizer_config.json", + "processor_config.json", + "chat_template.jinja", + "generation_config.json", + ]: + assert os.path.isfile(os.path.join(snapshot_dir, aux_file)) + assert os.path.isfile(os.path.join(snapshot_dir, "scheduler", SCHEDULER_CONFIG_NAME)) + # unrelated root files are still not downloaded + assert not os.path.isfile(os.path.join(snapshot_dir, "big_array.npy")) + def test_warning_unused_kwargs(self): model_id = "hf-internal-testing/unet-pipeline-dummy" logger = logging.get_logger("diffusers.pipelines") diff --git a/tests/pipelines/test_pipelines_auto.py b/tests/pipelines/test_pipelines_auto.py index f3c639c367f7..3d6c607b1a32 100644 --- a/tests/pipelines/test_pipelines_auto.py +++ b/tests/pipelines/test_pipelines_auto.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -31,9 +31,11 @@ DiffusionPipeline, ) from diffusers.pipelines.auto_pipeline import ( + AUTO_CONDITION2VIDEO_PIPELINES_MAPPING, AUTO_IMAGE2IMAGE_PIPELINES_MAPPING, AUTO_INPAINT_PIPELINES_MAPPING, AUTO_TEXT2IMAGE_PIPELINES_MAPPING, + _get_model, ) from ..testing_utils import slow @@ -50,6 +52,10 @@ class AutoPipelineFastTest(unittest.TestCase): + def test_condition2video_mapping(self): + assert AUTO_CONDITION2VIDEO_PIPELINES_MAPPING["cosmos3-omni"].__name__ == "Cosmos3OmniPipeline" + assert _get_model("Cosmos3OmniPipeline") == "cosmos3-omni" + @property def dummy_image_encoder(self): torch.manual_seed(0) diff --git a/tests/pipelines/test_pipelines_combined.py b/tests/pipelines/test_pipelines_combined.py index fffc053bae3f..2354decf4771 100644 --- a/tests/pipelines/test_pipelines_combined.py +++ b/tests/pipelines/test_pipelines_combined.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/pipelines/test_pipelines_common.py b/tests/pipelines/test_pipelines_common.py index d6a09028157f..139adac2b41e 100644 --- a/tests/pipelines/test_pipelines_common.py +++ b/tests/pipelines/test_pipelines_common.py @@ -137,7 +137,7 @@ def test_vae_slicing(self, image_count=4): output_1 = pipe(**inputs) # make sure sliced vae decode yields the same result - pipe.enable_vae_slicing() + pipe.vae.enable_slicing() inputs = self.get_dummy_inputs(device) inputs["prompt"] = [inputs["prompt"]] * image_count if "image" in inputs: @@ -164,7 +164,7 @@ def test_vae_tiling(self): output_1 = pipe(**inputs)[0] # make sure tiled vae decode yields the same result - pipe.enable_vae_tiling() + pipe.vae.enable_tiling() inputs = self.get_dummy_inputs(torch_device) inputs["return_dict"] = False output_2 = pipe(**inputs)[0] diff --git a/tests/pipelines/wan/test_wan_image_to_video.py b/tests/pipelines/wan/test_wan_image_to_video.py index 7ed263abdcb5..41cf2636d29f 100644 --- a/tests/pipelines/wan/test_wan_image_to_video.py +++ b/tests/pipelines/wan/test_wan_image_to_video.py @@ -154,7 +154,7 @@ def test_inference(self): self.assertEqual(generated_video.shape, (9, 3, 16, 16)) # fmt: off - expected_slice = torch.tensor([0.4525, 0.4525, 0.4497, 0.4536, 0.452, 0.4529, 0.454, 0.4535, 0.5072, 0.5527, 0.5165, 0.5244, 0.5481, 0.5282, 0.5208, 0.5214]) + expected_slice = torch.tensor([0.4528, 0.4525, 0.4493, 0.4537, 0.4521, 0.4532, 0.4543, 0.4536, 0.5084, 0.5252, 0.5211, 0.5120, 0.5419, 0.5355, 0.5169, 0.5213]) # fmt: on generated_slice = generated_video.flatten() @@ -329,7 +329,7 @@ def test_inference(self): self.assertEqual(generated_video.shape, (9, 3, 16, 16)) # fmt: off - expected_slice = torch.tensor([0.4531, 0.4527, 0.4498, 0.4542, 0.4526, 0.4527, 0.4534, 0.4534, 0.5061, 0.5185, 0.5283, 0.5181, 0.5309, 0.5365, 0.5113, 0.5244]) + expected_slice = torch.tensor([0.4525, 0.4525, 0.4497, 0.4537, 0.4520, 0.4529, 0.4540, 0.4535, 0.5157, 0.5449, 0.5201, 0.5192, 0.5398, 0.5374, 0.5162, 0.5112]) # fmt: on generated_slice = generated_video.flatten() diff --git a/tests/remote/test_remote_decode.py b/tests/remote/test_remote_decode.py index e48ddcb17438..9aee90e02a8b 100644 --- a/tests/remote/test_remote_decode.py +++ b/tests/remote/test_remote_decode.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/remote/test_remote_encode.py b/tests/remote/test_remote_encode.py index 4c0daf08fd8c..c186e73b8073 100644 --- a/tests/remote/test_remote_encode.py +++ b/tests/remote/test_remote_encode.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/schedulers/test_schedulers.py b/tests/schedulers/test_schedulers.py index 5a8380e659fc..914793716b79 100755 --- a/tests/schedulers/test_schedulers.py +++ b/tests/schedulers/test_schedulers.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/single_file/test_lumina2_transformer.py b/tests/single_file/test_lumina2_transformer.py index bb5a0bf473b6..51b8a2027a41 100644 --- a/tests/single_file/test_lumina2_transformer.py +++ b/tests/single_file/test_lumina2_transformer.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/single_file/test_model_autoencoder_dc_single_file.py b/tests/single_file/test_model_autoencoder_dc_single_file.py index 444ca4046977..cfedbe802cd6 100644 --- a/tests/single_file/test_model_autoencoder_dc_single_file.py +++ b/tests/single_file/test_model_autoencoder_dc_single_file.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/single_file/test_model_controlnet_single_file.py b/tests/single_file/test_model_controlnet_single_file.py index 2fa81fe3ae55..55c9feadd2c1 100644 --- a/tests/single_file/test_model_controlnet_single_file.py +++ b/tests/single_file/test_model_controlnet_single_file.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/single_file/test_model_flux_transformer_single_file.py b/tests/single_file/test_model_flux_transformer_single_file.py index bec886a02e45..6cd6b767d831 100644 --- a/tests/single_file/test_model_flux_transformer_single_file.py +++ b/tests/single_file/test_model_flux_transformer_single_file.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/single_file/test_model_motion_adapter_single_file.py b/tests/single_file/test_model_motion_adapter_single_file.py index a047c81b47aa..df2c687d11b7 100644 --- a/tests/single_file/test_model_motion_adapter_single_file.py +++ b/tests/single_file/test_model_motion_adapter_single_file.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/single_file/test_model_sd_cascade_unet_single_file.py b/tests/single_file/test_model_sd_cascade_unet_single_file.py index 7472122710eb..3d9e510aa8f7 100644 --- a/tests/single_file/test_model_sd_cascade_unet_single_file.py +++ b/tests/single_file/test_model_sd_cascade_unet_single_file.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/single_file/test_model_vae_single_file.py b/tests/single_file/test_model_vae_single_file.py index 9198d9b16337..ef0d885b108e 100644 --- a/tests/single_file/test_model_vae_single_file.py +++ b/tests/single_file/test_model_vae_single_file.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/single_file/test_model_wan_autoencoder_single_file.py b/tests/single_file/test_model_wan_autoencoder_single_file.py index 0babf302348f..728e57fd8868 100644 --- a/tests/single_file/test_model_wan_autoencoder_single_file.py +++ b/tests/single_file/test_model_wan_autoencoder_single_file.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tests/single_file/test_model_wan_transformer3d_single_file.py b/tests/single_file/test_model_wan_transformer3d_single_file.py index b76909206073..c7b140168565 100644 --- a/tests/single_file/test_model_wan_transformer3d_single_file.py +++ b/tests/single_file/test_model_wan_transformer3d_single_file.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/utils/generate_model_tests.py b/utils/generate_model_tests.py index cb54e50e4432..67ea5612ac12 100644 --- a/utils/generate_model_tests.py +++ b/utils/generate_model_tests.py @@ -1,6 +1,6 @@ #!/usr/bin/env python # coding=utf-8 -# Copyright 2025 HuggingFace Inc. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -437,7 +437,7 @@ def generate_test_file(model_info: dict, model_filepath: str, include_optional: lines = [ "# coding=utf-8", - "# Copyright 2025 HuggingFace Inc.", + "# Copyright 2026 HuggingFace Inc.", "#", '# Licensed under the Apache License, Version 2.0 (the "License");', "# you may not use this file except in compliance with the License.",