diff --git a/tests/lora/test_lora_layers_hunyuanvideo.py b/tests/lora/test_lora_layers_hunyuanvideo.py deleted file mode 100644 index 6bf64751fe7a..000000000000 --- a/tests/lora/test_lora_layers_hunyuanvideo.py +++ /dev/null @@ -1,242 +0,0 @@ -# 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. -# 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 gc -import sys -import unittest - -import numpy as np -import torch -from transformers import CLIPTextModel, CLIPTokenizer, LlamaModel, LlamaTokenizerFast - -from diffusers import ( - AutoencoderKLHunyuanVideo, - FlowMatchEulerDiscreteScheduler, - HunyuanVideoPipeline, - HunyuanVideoTransformer3DModel, -) - -from ..testing_utils import ( - Expectations, - backend_empty_cache, - floats_tensor, - nightly, - numpy_cosine_similarity_distance, - require_big_accelerator, - require_peft_backend, - require_torch_accelerator, - skip_mps, - torch_device, -) - - -sys.path.append(".") - -from .utils import PeftLoraLoaderMixinTests # noqa: E402 - - -@require_peft_backend -@skip_mps -class HunyuanVideoLoRATests(unittest.TestCase, PeftLoraLoaderMixinTests): - pipeline_class = HunyuanVideoPipeline - scheduler_cls = FlowMatchEulerDiscreteScheduler - scheduler_kwargs = {} - - transformer_kwargs = { - "in_channels": 4, - "out_channels": 4, - "num_attention_heads": 2, - "attention_head_dim": 10, - "num_layers": 1, - "num_single_layers": 1, - "num_refiner_layers": 1, - "patch_size": 1, - "patch_size_t": 1, - "guidance_embeds": True, - "text_embed_dim": 16, - "pooled_projection_dim": 8, - "rope_axes_dim": (2, 4, 4), - } - transformer_cls = HunyuanVideoTransformer3DModel - vae_kwargs = { - "in_channels": 3, - "out_channels": 3, - "latent_channels": 4, - "down_block_types": ( - "HunyuanVideoDownBlock3D", - "HunyuanVideoDownBlock3D", - "HunyuanVideoDownBlock3D", - "HunyuanVideoDownBlock3D", - ), - "up_block_types": ( - "HunyuanVideoUpBlock3D", - "HunyuanVideoUpBlock3D", - "HunyuanVideoUpBlock3D", - "HunyuanVideoUpBlock3D", - ), - "block_out_channels": (8, 8, 8, 8), - "layers_per_block": 1, - "act_fn": "silu", - "norm_num_groups": 4, - "scaling_factor": 0.476986, - "spatial_compression_ratio": 8, - "temporal_compression_ratio": 4, - "mid_block_add_attention": True, - } - vae_cls = AutoencoderKLHunyuanVideo - has_two_text_encoders = True - tokenizer_cls, tokenizer_id, tokenizer_subfolder = ( - LlamaTokenizerFast, - "hf-internal-testing/tiny-random-hunyuanvideo", - "tokenizer", - ) - tokenizer_2_cls, tokenizer_2_id, tokenizer_2_subfolder = ( - CLIPTokenizer, - "hf-internal-testing/tiny-random-hunyuanvideo", - "tokenizer_2", - ) - text_encoder_cls, text_encoder_id, text_encoder_subfolder = ( - LlamaModel, - "hf-internal-testing/tiny-random-hunyuanvideo", - "text_encoder", - ) - text_encoder_2_cls, text_encoder_2_id, text_encoder_2_subfolder = ( - CLIPTextModel, - "hf-internal-testing/tiny-random-hunyuanvideo", - "text_encoder_2", - ) - - supports_text_encoder_loras = False - - @property - def output_shape(self): - return (1, 9, 32, 32, 3) - - def get_dummy_inputs(self, with_generator=True): - batch_size = 1 - sequence_length = 16 - num_channels = 4 - num_frames = 9 - num_latent_frames = 3 # (num_frames - 1) // temporal_compression_ratio + 1 - sizes = (4, 4) - - generator = torch.manual_seed(0) - noise = floats_tensor((batch_size, num_latent_frames, num_channels) + sizes) - input_ids = torch.randint(1, sequence_length, size=(batch_size, sequence_length), generator=generator) - - pipeline_inputs = { - "prompt": "", - "num_frames": num_frames, - "num_inference_steps": 1, - "guidance_scale": 6.0, - "height": 32, - "width": 32, - "max_sequence_length": sequence_length, - "prompt_template": {"template": "{}", "crop_start": 0}, - "output_type": "np", - } - if with_generator: - pipeline_inputs.update({"generator": generator}) - - return noise, input_ids, pipeline_inputs - - def test_simple_inference_with_text_lora_denoiser_fused_multi(self): - super().test_simple_inference_with_text_lora_denoiser_fused_multi(expected_atol=9e-3) - - def test_simple_inference_with_text_denoiser_lora_unfused(self): - super().test_simple_inference_with_text_denoiser_lora_unfused(expected_atol=9e-3) - - # TODO(aryan): Fix the following test - @unittest.skip("This test fails with an error I haven't been able to debug yet.") - def test_simple_inference_save_pretrained(self): - pass - - @unittest.skip("Not supported in HunyuanVideo.") - def test_simple_inference_with_text_denoiser_block_scale(self): - pass - - @unittest.skip("Not supported in HunyuanVideo.") - def test_simple_inference_with_text_denoiser_block_scale_for_all_dict_options(self): - pass - - -@nightly -@require_torch_accelerator -@require_peft_backend -@require_big_accelerator -class HunyuanVideoLoRAIntegrationTests(unittest.TestCase): - """internal note: The integration slices were obtained on DGX. - - torch: 2.5.1+cu124 with CUDA 12.5. Need the same setup for the - assertions to pass. - """ - - num_inference_steps = 10 - seed = 0 - - def setUp(self): - super().setUp() - - gc.collect() - backend_empty_cache(torch_device) - - model_id = "hunyuanvideo-community/HunyuanVideo" - transformer = HunyuanVideoTransformer3DModel.from_pretrained( - model_id, subfolder="transformer", torch_dtype=torch.bfloat16 - ) - self.pipeline = HunyuanVideoPipeline.from_pretrained( - model_id, transformer=transformer, torch_dtype=torch.float16 - ).to(torch_device) - - def tearDown(self): - super().tearDown() - - gc.collect() - backend_empty_cache(torch_device) - - def test_original_format_cseti(self): - self.pipeline.load_lora_weights( - "Cseti/HunyuanVideo-LoRA-Arcane_Jinx-v1", weight_name="csetiarcane-nfjinx-v1-6000.safetensors" - ) - self.pipeline.fuse_lora() - self.pipeline.unload_lora_weights() - self.pipeline.vae.enable_tiling() - - prompt = "CSETIARCANE. A cat walks on the grass, realistic" - - out = self.pipeline( - prompt=prompt, - height=320, - width=512, - num_frames=9, - num_inference_steps=self.num_inference_steps, - output_type="np", - generator=torch.manual_seed(self.seed), - ).frames[0] - out = out.flatten() - out_slice = np.concatenate((out[:8], out[-8:])) - - # fmt: off - expected_slices = Expectations( - { - ("cuda", 7): np.array([0.1013, 0.1924, 0.0078, 0.1021, 0.1929, 0.0078, 0.1023, 0.1919, 0.7402, 0.104, 0.4482, 0.7354, 0.0925, 0.4382, 0.7275, 0.0815]), - ("xpu", 3): np.array([0.1013, 0.1924, 0.0078, 0.1021, 0.1929, 0.0078, 0.1023, 0.1919, 0.7402, 0.104, 0.4482, 0.7354, 0.0925, 0.4382, 0.7275, 0.0815]), - } - ) - # fmt: on - expected_slice = expected_slices.get_expectation() - - max_diff = numpy_cosine_similarity_distance(expected_slice.flatten(), out_slice) - - assert max_diff < 1e-3 diff --git a/tests/lora/test_lora_layers_ltx2.py b/tests/lora/test_lora_layers_ltx2.py deleted file mode 100644 index fab47cc3137e..000000000000 --- a/tests/lora/test_lora_layers_ltx2.py +++ /dev/null @@ -1,267 +0,0 @@ -# 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. -# 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 sys -import unittest - -import torch -from transformers import AutoTokenizer, Gemma3ForConditionalGeneration - -from diffusers import ( - AutoencoderKLLTX2Audio, - AutoencoderKLLTX2Video, - FlowMatchEulerDiscreteScheduler, - LTX2Pipeline, - LTX2VideoTransformer3DModel, -) -from diffusers.pipelines.ltx2 import LTX2TextConnectors -from diffusers.pipelines.ltx2.vocoder import LTX2Vocoder -from diffusers.utils.import_utils import is_peft_available - -from ..testing_utils import floats_tensor, require_peft_backend - - -if is_peft_available(): - from peft import LoraConfig - - -sys.path.append(".") - -from .utils import PeftLoraLoaderMixinTests # noqa: E402 - - -@require_peft_backend -class LTX2LoRATests(unittest.TestCase, PeftLoraLoaderMixinTests): - pipeline_class = LTX2Pipeline - scheduler_cls = FlowMatchEulerDiscreteScheduler - scheduler_kwargs = {} - - transformer_kwargs = { - "in_channels": 4, - "out_channels": 4, - "patch_size": 1, - "patch_size_t": 1, - "num_attention_heads": 2, - "attention_head_dim": 8, - "cross_attention_dim": 16, - "audio_in_channels": 4, - "audio_out_channels": 4, - "audio_num_attention_heads": 2, - "audio_attention_head_dim": 4, - "audio_cross_attention_dim": 8, - "num_layers": 1, - "qk_norm": "rms_norm_across_heads", - "caption_channels": 32, - "rope_double_precision": False, - "rope_type": "split", - } - transformer_cls = LTX2VideoTransformer3DModel - - vae_kwargs = { - "in_channels": 3, - "out_channels": 3, - "latent_channels": 4, - "block_out_channels": (8,), - "decoder_block_out_channels": (8,), - "layers_per_block": (1,), - "decoder_layers_per_block": (1, 1), - "spatio_temporal_scaling": (True,), - "decoder_spatio_temporal_scaling": (True,), - "decoder_inject_noise": (False, False), - "downsample_type": ("spatial",), - "upsample_residual": (False,), - "upsample_factor": (1,), - "timestep_conditioning": False, - "patch_size": 1, - "patch_size_t": 1, - "encoder_causal": True, - "decoder_causal": False, - } - vae_cls = AutoencoderKLLTX2Video - - audio_vae_kwargs = { - "base_channels": 4, - "output_channels": 2, - "ch_mult": (1,), - "num_res_blocks": 1, - "attn_resolutions": None, - "in_channels": 2, - "resolution": 32, - "latent_channels": 2, - "norm_type": "pixel", - "causality_axis": "height", - "dropout": 0.0, - "mid_block_add_attention": False, - "sample_rate": 16000, - "mel_hop_length": 160, - "is_causal": True, - "mel_bins": 8, - } - audio_vae_cls = AutoencoderKLLTX2Audio - - vocoder_kwargs = { - "in_channels": 16, # output_channels * mel_bins = 2 * 8 - "hidden_channels": 32, - "out_channels": 2, - "upsample_kernel_sizes": [4, 4], - "upsample_factors": [2, 2], - "resnet_kernel_sizes": [3], - "resnet_dilations": [[1, 3, 5]], - "leaky_relu_negative_slope": 0.1, - "output_sampling_rate": 16000, - } - vocoder_cls = LTX2Vocoder - - connectors_kwargs = { - "caption_channels": 32, # Will be set dynamically from text_encoder - "text_proj_in_factor": 2, # Will be set dynamically from text_encoder - "video_connector_num_attention_heads": 4, - "video_connector_attention_head_dim": 8, - "video_connector_num_layers": 1, - "video_connector_num_learnable_registers": None, - "audio_connector_num_attention_heads": 4, - "audio_connector_attention_head_dim": 8, - "audio_connector_num_layers": 1, - "audio_connector_num_learnable_registers": None, - "connector_rope_base_seq_len": 32, - "rope_theta": 10000.0, - "rope_double_precision": False, - "causal_temporal_positioning": False, - "rope_type": "split", - } - connectors_cls = LTX2TextConnectors - - tokenizer_cls, tokenizer_id = AutoTokenizer, "hf-internal-testing/tiny-gemma3" - text_encoder_cls, text_encoder_id = ( - Gemma3ForConditionalGeneration, - "hf-internal-testing/tiny-gemma3", - ) - - denoiser_target_modules = ["to_q", "to_k", "to_out.0"] - - supports_text_encoder_loras = False - - @property - def output_shape(self): - return (1, 5, 32, 32, 3) - - def get_dummy_inputs(self, with_generator=True): - batch_size = 1 - sequence_length = 16 - num_channels = 4 - num_frames = 5 - num_latent_frames = 2 - latent_height = 8 - latent_width = 8 - - generator = torch.manual_seed(0) - noise = floats_tensor((batch_size, num_latent_frames, num_channels, latent_height, latent_width)) - input_ids = torch.randint(1, sequence_length, size=(batch_size, sequence_length), generator=generator) - - pipeline_inputs = { - "prompt": "a robot dancing", - "num_frames": num_frames, - "num_inference_steps": 2, - "guidance_scale": 1.0, - "height": 32, - "width": 32, - "frame_rate": 25.0, - "max_sequence_length": sequence_length, - "output_type": "np", - } - if with_generator: - pipeline_inputs.update({"generator": generator}) - - return noise, input_ids, pipeline_inputs - - def get_dummy_components(self, scheduler_cls=None, use_dora=False, lora_alpha=None): - # Override to instantiate LTX2-specific components (connectors, audio_vae, vocoder) - torch.manual_seed(0) - text_encoder = self.text_encoder_cls.from_pretrained(self.text_encoder_id) - tokenizer = self.tokenizer_cls.from_pretrained(self.tokenizer_id) - - # Update caption_channels and text_proj_in_factor based on text_encoder config - transformer_kwargs = self.transformer_kwargs.copy() - transformer_kwargs["caption_channels"] = text_encoder.config.text_config.hidden_size - - connectors_kwargs = self.connectors_kwargs.copy() - connectors_kwargs["caption_channels"] = text_encoder.config.text_config.hidden_size - connectors_kwargs["text_proj_in_factor"] = text_encoder.config.text_config.num_hidden_layers + 1 - - torch.manual_seed(0) - transformer = self.transformer_cls(**transformer_kwargs) - - torch.manual_seed(0) - vae = self.vae_cls(**self.vae_kwargs) - vae.use_framewise_encoding = False - vae.use_framewise_decoding = False - - torch.manual_seed(0) - audio_vae = self.audio_vae_cls(**self.audio_vae_kwargs) - - torch.manual_seed(0) - vocoder = self.vocoder_cls(**self.vocoder_kwargs) - - torch.manual_seed(0) - connectors = self.connectors_cls(**connectors_kwargs) - - if scheduler_cls is None: - scheduler_cls = self.scheduler_cls - scheduler = scheduler_cls(**self.scheduler_kwargs) - - rank = 4 - lora_alpha = rank if lora_alpha is None else lora_alpha - - text_lora_config = LoraConfig( - r=rank, - lora_alpha=lora_alpha, - target_modules=self.text_encoder_target_modules, - init_lora_weights=False, - use_dora=use_dora, - ) - - denoiser_lora_config = LoraConfig( - r=rank, - lora_alpha=lora_alpha, - target_modules=["to_q", "to_k", "to_v", "to_out.0"], - init_lora_weights=False, - use_dora=use_dora, - ) - - pipeline_components = { - "transformer": transformer, - "vae": vae, - "audio_vae": audio_vae, - "scheduler": scheduler, - "text_encoder": text_encoder, - "tokenizer": tokenizer, - "connectors": connectors, - "vocoder": vocoder, - } - - return pipeline_components, text_lora_config, denoiser_lora_config - - def test_simple_inference_with_text_lora_denoiser_fused_multi(self): - super().test_simple_inference_with_text_lora_denoiser_fused_multi(expected_atol=9e-3) - - def test_simple_inference_with_text_denoiser_lora_unfused(self): - super().test_simple_inference_with_text_denoiser_lora_unfused(expected_atol=9e-3) - - @unittest.skip("Not supported in LTX2.") - def test_simple_inference_with_text_denoiser_block_scale(self): - pass - - @unittest.skip("Not supported in LTX2.") - def test_simple_inference_with_text_denoiser_block_scale_for_all_dict_options(self): - pass diff --git a/tests/lora/test_lora_layers_ltx_video.py b/tests/lora/test_lora_layers_ltx_video.py deleted file mode 100644 index d1ae1ba1f468..000000000000 --- a/tests/lora/test_lora_layers_ltx_video.py +++ /dev/null @@ -1,125 +0,0 @@ -# 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. -# 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 sys -import unittest - -import torch -from transformers import AutoTokenizer, T5EncoderModel - -from diffusers import ( - AutoencoderKLLTXVideo, - FlowMatchEulerDiscreteScheduler, - LTXPipeline, - LTXVideoTransformer3DModel, -) - -from ..testing_utils import floats_tensor, require_peft_backend - - -sys.path.append(".") - -from .utils import PeftLoraLoaderMixinTests # noqa: E402 - - -@require_peft_backend -class LTXVideoLoRATests(unittest.TestCase, PeftLoraLoaderMixinTests): - pipeline_class = LTXPipeline - scheduler_cls = FlowMatchEulerDiscreteScheduler - scheduler_kwargs = {} - - transformer_kwargs = { - "in_channels": 8, - "out_channels": 8, - "patch_size": 1, - "patch_size_t": 1, - "num_attention_heads": 4, - "attention_head_dim": 8, - "cross_attention_dim": 32, - "num_layers": 1, - "caption_channels": 32, - } - transformer_cls = LTXVideoTransformer3DModel - vae_kwargs = { - "in_channels": 3, - "out_channels": 3, - "latent_channels": 8, - "block_out_channels": (8, 8, 8, 8), - "decoder_block_out_channels": (8, 8, 8, 8), - "layers_per_block": (1, 1, 1, 1, 1), - "decoder_layers_per_block": (1, 1, 1, 1, 1), - "spatio_temporal_scaling": (True, True, False, False), - "decoder_spatio_temporal_scaling": (True, True, False, False), - "decoder_inject_noise": (False, False, False, False, False), - "upsample_residual": (False, False, False, False), - "upsample_factor": (1, 1, 1, 1), - "timestep_conditioning": False, - "patch_size": 1, - "patch_size_t": 1, - "encoder_causal": True, - "decoder_causal": False, - } - vae_cls = AutoencoderKLLTXVideo - tokenizer_cls, tokenizer_id = AutoTokenizer, "hf-internal-testing/tiny-random-t5" - text_encoder_cls, text_encoder_id = T5EncoderModel, "hf-internal-testing/tiny-random-t5" - - text_encoder_target_modules = ["q", "k", "v", "o"] - - supports_text_encoder_loras = False - - @property - def output_shape(self): - return (1, 9, 32, 32, 3) - - def get_dummy_inputs(self, with_generator=True): - batch_size = 1 - sequence_length = 16 - num_channels = 8 - num_frames = 9 - num_latent_frames = 3 # (num_frames - 1) // temporal_compression_ratio + 1 - latent_height = 8 - latent_width = 8 - - generator = torch.manual_seed(0) - noise = floats_tensor((batch_size, num_latent_frames, num_channels, latent_height, latent_width)) - input_ids = torch.randint(1, sequence_length, size=(batch_size, sequence_length), generator=generator) - - pipeline_inputs = { - "prompt": "dance monkey", - "num_frames": num_frames, - "num_inference_steps": 4, - "guidance_scale": 6.0, - "height": 32, - "width": 32, - "max_sequence_length": sequence_length, - "output_type": "np", - } - if with_generator: - pipeline_inputs.update({"generator": generator}) - - return noise, input_ids, pipeline_inputs - - def test_simple_inference_with_text_lora_denoiser_fused_multi(self): - super().test_simple_inference_with_text_lora_denoiser_fused_multi(expected_atol=9e-3) - - def test_simple_inference_with_text_denoiser_lora_unfused(self): - super().test_simple_inference_with_text_denoiser_lora_unfused(expected_atol=9e-3) - - @unittest.skip("Not supported in LTXVideo.") - def test_simple_inference_with_text_denoiser_block_scale(self): - pass - - @unittest.skip("Not supported in LTXVideo.") - def test_simple_inference_with_text_denoiser_block_scale_for_all_dict_options(self): - pass diff --git a/tests/pipelines/hunyuan_video/test_hunyuan_video.py b/tests/pipelines/hunyuan_video/test_hunyuan_video.py index b16c8d7158a6..6b857eecea3b 100644 --- a/tests/pipelines/hunyuan_video/test_hunyuan_video.py +++ b/tests/pipelines/hunyuan_video/test_hunyuan_video.py @@ -12,68 +12,58 @@ # See the License for the specific language governing permissions and # limitations under the License. -import inspect -import unittest +import gc import numpy as np +import pytest import torch from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer, LlamaConfig, LlamaModel, LlamaTokenizer from diffusers import ( AutoencoderKLHunyuanVideo, - FasterCacheConfig, FlowMatchEulerDiscreteScheduler, HunyuanVideoPipeline, HunyuanVideoTransformer3DModel, ) -from ...testing_utils import enable_full_determinism, torch_device -from ..test_pipelines_common import ( +from ...testing_utils import ( + Expectations, + assert_tensors_close, + backend_empty_cache, + enable_full_determinism, + nightly, + numpy_cosine_similarity_distance, + require_big_accelerator, + require_peft_backend, + require_torch_accelerator, + torch_device, +) +from ..testing_utils import ( + BasePipelineTesterConfig, FasterCacheTesterMixin, FirstBlockCacheTesterMixin, + LoraMemoryTesterMixin, + LoraTesterMixin, + MemoryTesterMixin, PipelineTesterMixin, PyramidAttentionBroadcastTesterMixin, TaylorSeerCacheTesterMixin, - to_np, ) enable_full_determinism() -class HunyuanVideoPipelineFastTests( - PipelineTesterMixin, - PyramidAttentionBroadcastTesterMixin, - FasterCacheTesterMixin, - FirstBlockCacheTesterMixin, - TaylorSeerCacheTesterMixin, - unittest.TestCase, -): +class HunyuanVideoPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = HunyuanVideoPipeline - params = frozenset(["prompt", "height", "width", "guidance_scale", "prompt_embeds", "pooled_prompt_embeds"]) - batch_params = frozenset(["prompt"]) - required_optional_params = frozenset( - [ - "num_inference_steps", - "generator", - "latents", - "return_dict", - "callback_on_step_end", - "callback_on_step_end_tensor_inputs", - ] + required_input_params_in_call_signature = frozenset( + ["prompt", "height", "width", "guidance_scale", "prompt_embeds", "pooled_prompt_embeds"] ) - - # there is no xformers processor for Flux - test_xformers_attention = False - test_layerwise_casting = True - test_group_offloading = True - - faster_cache_config = FasterCacheConfig( - spatial_attention_block_skip_range=2, - spatial_attention_timestep_skip_range=(-1, 901), - unconditional_batch_skip_range=2, - attention_weight_callback=lambda _: 0.5, - is_guidance_distilled=True, + batch_input_params = frozenset(["prompt"]) + output_shape = (9, 3, 16, 16) + # HunyuanVideo is a video pipeline: it exposes `num_videos_per_prompt`, not `num_images_per_prompt`. + optional_input_params = frozenset( + ["num_inference_steps", "num_videos_per_prompt", "generator", "latents", "output_type", "return_dict"] ) def get_dummy_components(self, num_layers: int = 1, num_single_layers: int = 1): @@ -159,7 +149,7 @@ def get_dummy_components(self, num_layers: int = 1, num_single_layers: int = 1): text_encoder_2 = CLIPTextModel(clip_text_encoder_config) tokenizer_2 = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") - components = { + return { "transformer": transformer, "vae": vae, "scheduler": scheduler, @@ -168,21 +158,15 @@ def get_dummy_components(self, num_layers: int = 1, num_single_layers: int = 1): "tokenizer": tokenizer, "tokenizer_2": tokenizer_2, } - 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 = { + def get_dummy_inputs(self): + return { "prompt": "dance monkey", "prompt_template": { "template": "{}", "crop_start": 0, }, - "generator": generator, + "generator": self.get_generator(0), "num_inference_steps": 2, "guidance_scale": 4.5, "height": 16, @@ -190,22 +174,20 @@ def get_dummy_inputs(self, device, seed=0): # 4 * k + 1 is the recommendation "num_frames": 9, "max_sequence_length": 16, + # Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`). "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) +class TestHunyuanVideoPipeline(HunyuanVideoPipelineTesterConfig, PipelineTesterMixin): + def test_inference(self): + # Run on CPU: the expected slice below is CPU-specific. + pipe = self.get_pipeline() - inputs = self.get_dummy_inputs(device) + inputs = self.get_dummy_inputs() video = pipe(**inputs).frames generated_video = video[0] - self.assertEqual(generated_video.shape, (9, 3, 16, 16)) + assert generated_video.shape == self.output_shape # fmt: off 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]) @@ -213,117 +195,16 @@ def test_inference(self): generated_slice = generated_video.flatten() generated_slice = torch.cat([generated_slice[:8], generated_slice[-8:]]) - self.assertTrue( - torch.allclose(generated_slice, expected_slice, atol=1e-3), - "The generated video does not match the expected slice.", + assert_tensors_close( + generated_slice, expected_slice, atol=1e-3, msg="The generated video does not match the expected slice." ) - def test_callback_inputs(self): - sig = inspect.signature(self.pipeline_class.__call__) - has_callback_tensor_inputs = "callback_on_step_end_tensor_inputs" in sig.parameters - has_callback_step_end = "callback_on_step_end" in sig.parameters - - if not (has_callback_tensor_inputs and has_callback_step_end): - return - - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - self.assertTrue( - hasattr(pipe, "_callback_tensor_inputs"), - f" {self.pipeline_class} should have `_callback_tensor_inputs` that defines a list of tensor variables its callback function can use as inputs", - ) - - def callback_inputs_subset(pipe, i, t, callback_kwargs): - # iterate over callback args - for tensor_name, tensor_value in callback_kwargs.items(): - # check that we're only passing in allowed tensor inputs - assert tensor_name in pipe._callback_tensor_inputs - - return callback_kwargs - - def callback_inputs_all(pipe, i, t, callback_kwargs): - for tensor_name in pipe._callback_tensor_inputs: - assert tensor_name in callback_kwargs - - # iterate over callback args - for tensor_name, tensor_value in callback_kwargs.items(): - # check that we're only passing in allowed tensor inputs - assert tensor_name in pipe._callback_tensor_inputs - - return callback_kwargs - - inputs = self.get_dummy_inputs(torch_device) - - # Test passing in a subset - inputs["callback_on_step_end"] = callback_inputs_subset - inputs["callback_on_step_end_tensor_inputs"] = ["latents"] - output = pipe(**inputs)[0] - - # Test passing in a everything - inputs["callback_on_step_end"] = callback_inputs_all - inputs["callback_on_step_end_tensor_inputs"] = pipe._callback_tensor_inputs - output = pipe(**inputs)[0] - - def callback_inputs_change_tensor(pipe, i, t, callback_kwargs): - is_last = i == (pipe.num_timesteps - 1) - if is_last: - callback_kwargs["latents"] = torch.zeros_like(callback_kwargs["latents"]) - return callback_kwargs - - inputs["callback_on_step_end"] = callback_inputs_change_tensor - inputs["callback_on_step_end_tensor_inputs"] = pipe._callback_tensor_inputs - output = pipe(**inputs)[0] - assert output.abs().sum() < 1e10 - - def test_attention_slicing_forward_pass( - self, test_max_difference=True, test_mean_pixel_difference=True, expected_max_diff=1e-3 - ): - if not self.test_attention_slicing: - return - - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - for component in pipe.components.values(): - if hasattr(component, "set_default_attn_processor"): - component.set_default_attn_processor() - pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - - generator_device = "cpu" - inputs = self.get_dummy_inputs(generator_device) - output_without_slicing = pipe(**inputs)[0] - - pipe.enable_attention_slicing(slice_size=1) - inputs = self.get_dummy_inputs(generator_device) - output_with_slicing1 = pipe(**inputs)[0] - - pipe.enable_attention_slicing(slice_size=2) - inputs = self.get_dummy_inputs(generator_device) - output_with_slicing2 = pipe(**inputs)[0] - - if test_max_difference: - max_diff1 = np.abs(to_np(output_with_slicing1) - to_np(output_without_slicing)).max() - max_diff2 = np.abs(to_np(output_with_slicing2) - to_np(output_without_slicing)).max() - self.assertLess( - max(max_diff1, max_diff2), - expected_max_diff, - "Attention slicing should not affect the inference results", - ) - - def test_vae_tiling(self, expected_diff_max: float = 0.2): + def test_vae_tiling(self, expected_diff_max: float = 0.6): # Seems to require higher tolerance than the other tests - expected_diff_max = 0.6 - generator_device = "cpu" - components = self.get_dummy_components() - - pipe = self.pipeline_class(**components) - pipe.to("cpu") - pipe.set_progress_bar_config(disable=None) + pipe = self.get_pipeline().to(torch_device) # Without tiling - inputs = self.get_dummy_inputs(generator_device) + inputs = self.get_dummy_inputs() inputs["height"] = inputs["width"] = 128 output_without_tiling = pipe(**inputs)[0] @@ -334,25 +215,131 @@ def test_vae_tiling(self, expected_diff_max: float = 0.2): tile_sample_stride_height=64, tile_sample_stride_width=64, ) - inputs = self.get_dummy_inputs(generator_device) + inputs = self.get_dummy_inputs() inputs["height"] = inputs["width"] = 128 output_with_tiling = pipe(**inputs)[0] - self.assertLess( - (to_np(output_without_tiling) - to_np(output_with_tiling)).max(), - expected_diff_max, - "VAE tiling should not affect the inference results", + assert (output_without_tiling - output_with_tiling).abs().max() < expected_diff_max, ( + "VAE tiling should not affect the inference results." ) # TODO(aryan): Create a dummy gemma model with smol vocab size - @unittest.skip( + @pytest.mark.skip( "A very small vocab size is used for fast tests. So, Any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error." ) def test_inference_batch_consistent(self): pass - @unittest.skip( + @pytest.mark.skip( "A very small vocab size is used for fast tests. So, Any kind of prompt other than the empty default used in other tests will lead to a embedding lookup error. This test uses a long prompt that causes the error." ) def test_inference_batch_single_identical(self): pass + + +class TestHunyuanVideoPipelineMemory(HunyuanVideoPipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the HunyuanVideo pipeline.""" + + +class TestHunyuanVideoPipelinePyramidAttentionBroadcast( + HunyuanVideoPipelineTesterConfig, PyramidAttentionBroadcastTesterMixin +): + """Pyramid Attention Broadcast cache tests for the HunyuanVideo pipeline.""" + + +class TestHunyuanVideoPipelineFasterCache(HunyuanVideoPipelineTesterConfig, FasterCacheTesterMixin): + """FasterCache tests for the HunyuanVideo pipeline.""" + + # HunyuanVideo is guidance-distilled, so the FasterCache tester must skip the low/high-frequency-delta checks. + FASTER_CACHE_CONFIG = { + "spatial_attention_block_skip_range": 2, + "spatial_attention_timestep_skip_range": (-1, 901), + "unconditional_batch_skip_range": 2, + "attention_weight_callback": lambda _: 0.5, + "is_guidance_distilled": True, + } + + +class TestHunyuanVideoPipelineFirstBlockCache(HunyuanVideoPipelineTesterConfig, FirstBlockCacheTesterMixin): + """First Block Cache tests for the HunyuanVideo pipeline.""" + + +class TestHunyuanVideoPipelineTaylorSeerCache(HunyuanVideoPipelineTesterConfig, TaylorSeerCacheTesterMixin): + """TaylorSeer cache tests for the HunyuanVideo pipeline.""" + + +class TestHunyuanVideoPipelineLoRA(HunyuanVideoPipelineTesterConfig, LoraTesterMixin): + """LoRA tests for the HunyuanVideo pipeline.""" + + +class TestHunyuanVideoPipelineLoRAMemory(HunyuanVideoPipelineTesterConfig, LoraMemoryTesterMixin): + """LoRA x memory-optimization tests (group offload, CPU offload) for the HunyuanVideo pipeline.""" + + +@nightly +@require_torch_accelerator +@require_peft_backend +@require_big_accelerator +class TestHunyuanVideoLoRAIntegration: + """internal note: The integration slices were obtained on DGX. + + torch: 2.5.1+cu124 with CUDA 12.5. Need the same setup for the + assertions to pass. + """ + + num_inference_steps = 10 + seed = 0 + repo_id = "hunyuanvideo-community/HunyuanVideo" + + @pytest.fixture(autouse=True) + def cleanup(self): + gc.collect() + backend_empty_cache(torch_device) + yield + gc.collect() + backend_empty_cache(torch_device) + + @pytest.fixture + def pipeline(self): + transformer = HunyuanVideoTransformer3DModel.from_pretrained( + self.repo_id, subfolder="transformer", dtype=torch.bfloat16 + ) + return HunyuanVideoPipeline.from_pretrained(self.repo_id, transformer=transformer, dtype=torch.float16).to( + torch_device + ) + + def test_original_format_cseti(self, pipeline): + pipeline.load_lora_weights( + "Cseti/HunyuanVideo-LoRA-Arcane_Jinx-v1", weight_name="csetiarcane-nfjinx-v1-6000.safetensors" + ) + pipeline.fuse_lora() + pipeline.unload_lora_weights() + pipeline.vae.enable_tiling() + + prompt = "CSETIARCANE. A cat walks on the grass, realistic" + + out = pipeline( + prompt=prompt, + height=320, + width=512, + num_frames=9, + num_inference_steps=self.num_inference_steps, + output_type="np", + generator=torch.manual_seed(self.seed), + ).frames[0] + out = out.flatten() + out_slice = np.concatenate((out[:8], out[-8:])) + + # fmt: off + expected_slices = Expectations( + { + ("cuda", 7): np.array([0.1013, 0.1924, 0.0078, 0.1021, 0.1929, 0.0078, 0.1023, 0.1919, 0.7402, 0.104, 0.4482, 0.7354, 0.0925, 0.4382, 0.7275, 0.0815]), + ("xpu", 3): np.array([0.1013, 0.1924, 0.0078, 0.1021, 0.1929, 0.0078, 0.1023, 0.1919, 0.7402, 0.104, 0.4482, 0.7354, 0.0925, 0.4382, 0.7275, 0.0815]), + } + ) + # fmt: on + expected_slice = expected_slices.get_expectation() + + max_diff = numpy_cosine_similarity_distance(expected_slice.flatten(), out_slice) + + assert max_diff < 1e-3 diff --git a/tests/pipelines/ltx/test_ltx.py b/tests/pipelines/ltx/test_ltx.py index da74c77dadc3..79d7be65f8aa 100644 --- a/tests/pipelines/ltx/test_ltx.py +++ b/tests/pipelines/ltx/test_ltx.py @@ -12,42 +12,36 @@ # See the License for the specific language governing permissions and # limitations under the License. -import inspect -import unittest - -import numpy as np import torch from transformers import AutoConfig, AutoTokenizer, T5EncoderModel from diffusers import AutoencoderKLLTXVideo, FlowMatchEulerDiscreteScheduler, LTXPipeline, LTXVideoTransformer3DModel from ...testing_utils import enable_full_determinism, torch_device -from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS -from ..test_pipelines_common import FirstBlockCacheTesterMixin, PipelineTesterMixin, to_np +from ..testing_utils import ( + BasePipelineTesterConfig, + FirstBlockCacheTesterMixin, + LoraMemoryTesterMixin, + LoraTesterMixin, + MemoryTesterMixin, + PipelineTesterMixin, +) enable_full_determinism() -class LTXPipelineFastTests(PipelineTesterMixin, FirstBlockCacheTesterMixin, unittest.TestCase): +class LTXPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = LTXPipeline - params = TEXT_TO_IMAGE_PARAMS - {"cross_attention_kwargs"} - batch_params = TEXT_TO_IMAGE_BATCH_PARAMS - image_params = TEXT_TO_IMAGE_IMAGE_PARAMS - image_latents_params = TEXT_TO_IMAGE_IMAGE_PARAMS - required_optional_params = frozenset( - [ - "num_inference_steps", - "generator", - "latents", - "return_dict", - "callback_on_step_end", - "callback_on_step_end_tensor_inputs", - ] + required_input_params_in_call_signature = frozenset( + ["prompt", "height", "width", "guidance_scale", "negative_prompt", "prompt_embeds", "negative_prompt_embeds"] + ) + batch_input_params = frozenset(["prompt", "negative_prompt"]) + output_shape = (9, 3, 32, 32) + # LTX is a video pipeline: it exposes `num_videos_per_prompt`, not the base default `num_images_per_prompt`. + optional_input_params = frozenset( + ["num_inference_steps", "num_videos_per_prompt", "generator", "latents", "output_type", "return_dict"] ) - test_xformers_attention = False - test_layerwise_casting = True - test_group_offloading = True def get_dummy_components(self, num_layers: int = 1): torch.manual_seed(0) @@ -94,25 +88,19 @@ def get_dummy_components(self, num_layers: int = 1): text_encoder = T5EncoderModel(config).eval() tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") - components = { + return { "transformer": transformer, "vae": vae, "scheduler": scheduler, "text_encoder": text_encoder, "tokenizer": tokenizer, } - 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 = { + def get_dummy_inputs(self): + return { "prompt": "dance monkey", "negative_prompt": "", - "generator": generator, + "generator": self.get_generator(0), "num_inference_steps": 2, "guidance_scale": 3.0, "height": 32, @@ -120,135 +108,20 @@ def get_dummy_inputs(self, device, seed=0): # 8 * k + 1 is the recommendation "num_frames": 9, "max_sequence_length": 16, + # Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`). "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) - video = pipe(**inputs).frames - generated_video = video[0] - - self.assertEqual(generated_video.shape, (9, 3, 32, 32)) - expected_video = torch.randn(9, 3, 32, 32) - max_diff = np.abs(generated_video - expected_video).max() - self.assertLessEqual(max_diff, 1e10) - - def test_callback_inputs(self): - sig = inspect.signature(self.pipeline_class.__call__) - has_callback_tensor_inputs = "callback_on_step_end_tensor_inputs" in sig.parameters - has_callback_step_end = "callback_on_step_end" in sig.parameters - - if not (has_callback_tensor_inputs and has_callback_step_end): - return - - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - self.assertTrue( - hasattr(pipe, "_callback_tensor_inputs"), - f" {self.pipeline_class} should have `_callback_tensor_inputs` that defines a list of tensor variables its callback function can use as inputs", - ) - - def callback_inputs_subset(pipe, i, t, callback_kwargs): - # iterate over callback args - for tensor_name, tensor_value in callback_kwargs.items(): - # check that we're only passing in allowed tensor inputs - assert tensor_name in pipe._callback_tensor_inputs - - return callback_kwargs - - def callback_inputs_all(pipe, i, t, callback_kwargs): - for tensor_name in pipe._callback_tensor_inputs: - assert tensor_name in callback_kwargs - - # iterate over callback args - for tensor_name, tensor_value in callback_kwargs.items(): - # check that we're only passing in allowed tensor inputs - assert tensor_name in pipe._callback_tensor_inputs - - return callback_kwargs - - inputs = self.get_dummy_inputs(torch_device) - - # Test passing in a subset - inputs["callback_on_step_end"] = callback_inputs_subset - inputs["callback_on_step_end_tensor_inputs"] = ["latents"] - output = pipe(**inputs)[0] - # Test passing in a everything - inputs["callback_on_step_end"] = callback_inputs_all - inputs["callback_on_step_end_tensor_inputs"] = pipe._callback_tensor_inputs - output = pipe(**inputs)[0] - - def callback_inputs_change_tensor(pipe, i, t, callback_kwargs): - is_last = i == (pipe.num_timesteps - 1) - if is_last: - callback_kwargs["latents"] = torch.zeros_like(callback_kwargs["latents"]) - return callback_kwargs - - inputs["callback_on_step_end"] = callback_inputs_change_tensor - inputs["callback_on_step_end_tensor_inputs"] = pipe._callback_tensor_inputs - output = pipe(**inputs)[0] - assert output.abs().sum() < 1e10 - - def test_inference_batch_single_identical(self): - self._test_inference_batch_single_identical(batch_size=3, expected_max_diff=1e-3) - - def test_attention_slicing_forward_pass( - self, test_max_difference=True, test_mean_pixel_difference=True, expected_max_diff=1e-3 - ): - if not self.test_attention_slicing: - return - - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - for component in pipe.components.values(): - if hasattr(component, "set_default_attn_processor"): - component.set_default_attn_processor() - pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - - generator_device = "cpu" - inputs = self.get_dummy_inputs(generator_device) - output_without_slicing = pipe(**inputs)[0] - - pipe.enable_attention_slicing(slice_size=1) - inputs = self.get_dummy_inputs(generator_device) - output_with_slicing1 = pipe(**inputs)[0] - - pipe.enable_attention_slicing(slice_size=2) - inputs = self.get_dummy_inputs(generator_device) - output_with_slicing2 = pipe(**inputs)[0] - - if test_max_difference: - max_diff1 = np.abs(to_np(output_with_slicing1) - to_np(output_without_slicing)).max() - max_diff2 = np.abs(to_np(output_with_slicing2) - to_np(output_without_slicing)).max() - self.assertLess( - max(max_diff1, max_diff2), - expected_max_diff, - "Attention slicing should not affect the inference results", - ) +class TestLTXPipeline(LTXPipelineTesterConfig, PipelineTesterMixin): + def test_inference_batch_single_identical(self, batch_size=3, expected_max_diff=1e-3): + super().test_inference_batch_single_identical(batch_size=batch_size, expected_max_diff=expected_max_diff) def test_vae_tiling(self, expected_diff_max: float = 0.2): - generator_device = "cpu" - components = self.get_dummy_components() - - pipe = self.pipeline_class(**components) - pipe.to("cpu") - pipe.set_progress_bar_config(disable=None) + pipe = self.get_pipeline().to(torch_device) # Without tiling - inputs = self.get_dummy_inputs(generator_device) + inputs = self.get_dummy_inputs() inputs["height"] = inputs["width"] = 128 output_without_tiling = pipe(**inputs)[0] @@ -259,12 +132,26 @@ def test_vae_tiling(self, expected_diff_max: float = 0.2): tile_sample_stride_height=64, tile_sample_stride_width=64, ) - inputs = self.get_dummy_inputs(generator_device) + inputs = self.get_dummy_inputs() inputs["height"] = inputs["width"] = 128 output_with_tiling = pipe(**inputs)[0] - self.assertLess( - (to_np(output_without_tiling) - to_np(output_with_tiling)).max(), - expected_diff_max, - "VAE tiling should not affect the inference results", + assert (output_without_tiling - output_with_tiling).abs().max() < expected_diff_max, ( + "VAE tiling should not affect the inference results." ) + + +class TestLTXPipelineMemory(LTXPipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the LTX pipeline.""" + + +class TestLTXPipelineFirstBlockCache(LTXPipelineTesterConfig, FirstBlockCacheTesterMixin): + """First-block-cache tests for the LTX pipeline.""" + + +class TestLTXPipelineLoRA(LTXPipelineTesterConfig, LoraTesterMixin): + """LoRA tests for the LTX pipeline.""" + + +class TestLTXPipelineLoRAMemory(LTXPipelineTesterConfig, LoraMemoryTesterMixin): + """LoRA x memory-optimization tests (group offload, CPU offload) for the LTX pipeline.""" diff --git a/tests/pipelines/ltx2/test_ltx2.py b/tests/pipelines/ltx2/test_ltx2.py index 917c4b1d6d88..89b7724b4351 100644 --- a/tests/pipelines/ltx2/test_ltx2.py +++ b/tests/pipelines/ltx2/test_ltx2.py @@ -12,8 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -import unittest - +import pytest import torch from transformers import AutoTokenizer, Gemma3ForConditionalGeneration @@ -27,34 +26,39 @@ from diffusers.pipelines.ltx2 import LTX2DurationHead, LTX2TextConnectors from diffusers.pipelines.ltx2.vocoder import LTX2Vocoder -from ...testing_utils import enable_full_determinism, torch_device -from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS -from ..test_pipelines_common import PipelineTesterMixin +from ...testing_utils import assert_tensors_close, enable_full_determinism, require_torch_accelerator, torch_device +from ..testing_utils import ( + BasePipelineTesterConfig, + LoraMemoryTesterMixin, + LoraTesterMixin, + MemoryTesterMixin, + PipelineTesterMixin, +) enable_full_determinism() -class LTX2PipelineFastTests(PipelineTesterMixin, unittest.TestCase): +class LTX2PipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = LTX2Pipeline - params = TEXT_TO_IMAGE_PARAMS - {"cross_attention_kwargs"} - batch_params = TEXT_TO_IMAGE_BATCH_PARAMS - image_params = TEXT_TO_IMAGE_IMAGE_PARAMS - image_latents_params = TEXT_TO_IMAGE_IMAGE_PARAMS - required_optional_params = frozenset( + required_input_params_in_call_signature = frozenset( + ["prompt", "height", "width", "guidance_scale", "negative_prompt", "prompt_embeds", "negative_prompt_embeds"] + ) + batch_input_params = frozenset(["prompt", "negative_prompt"]) + output_shape = (5, 3, 32, 32) + # LTX2 is a video pipeline (`num_videos_per_prompt`, not `num_images_per_prompt`) and takes a second latent + # input for the audio stream. + optional_input_params = frozenset( [ "num_inference_steps", + "num_videos_per_prompt", "generator", "latents", "audio_latents", "output_type", "return_dict", - "callback_on_step_end", - "callback_on_step_end_tensor_inputs", ] ) - test_attention_slicing = False - test_xformers_attention = False base_text_encoder_ckpt_id = "hf-internal-testing/tiny-gemma3" @@ -161,7 +165,7 @@ def get_dummy_components(self): scheduler = FlowMatchEulerDiscreteScheduler() - components = { + return { "transformer": transformer, "vae": vae, "audio_vae": audio_vae, @@ -175,18 +179,11 @@ def get_dummy_components(self): "duration_head": None, } - 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 = { + def get_dummy_inputs(self): + return { "prompt": "a robot dancing", "negative_prompt": "", - "generator": generator, + "generator": self.get_generator(0), "num_inference_steps": 2, "guidance_scale": 1.0, # Pin legacy sampling knobs so deterministic slice tests stay stable when @@ -205,27 +202,41 @@ def get_dummy_inputs(self, device, seed=0): "num_frames": 5, "frame_rate": 25.0, "max_sequence_length": 16, + # Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`). "output_type": "pt", } - return inputs - - def test_inference(self): - device = "cpu" + def get_dummy_duration_head(self): + torch.manual_seed(0) + # The dummy connectors emit 4 heads * 8 head_dim = 32 wide output for both streams. + return LTX2DurationHead( + video_cross_attention_dim=32, + audio_cross_attention_dim=32, + pooler_hidden_dim=8, + num_queries=1, + num_pooler_heads=2, + mlp_hidden_dim=8, + ) + def get_pipeline_with_duration_head(self): components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe.to(device) - pipe.set_progress_bar_config(disable=None) + components["duration_head"] = self.get_dummy_duration_head() + return self.get_pipeline(**components).to(torch_device) + + +class TestLTX2Pipeline(LTX2PipelineTesterConfig, PipelineTesterMixin): + def test_inference(self): + # Run on CPU: the expected slices below are CPU-specific. + pipe = self.get_pipeline() - inputs = self.get_dummy_inputs(device) + inputs = self.get_dummy_inputs() output = pipe(**inputs) video = output.frames audio = output.audio - self.assertEqual(video.shape, (1, 5, 3, 32, 32)) - self.assertEqual(audio.shape[0], 1) - self.assertEqual(audio.shape[1], components["vocoder"].config.out_channels) + assert video.shape == (1, *self.output_shape) + assert audio.shape[0] == 1 + assert audio.shape[1] == pipe.vocoder.config.out_channels # fmt: off expected_video_slice = torch.tensor( @@ -245,26 +256,22 @@ def test_inference(self): generated_video_slice = torch.cat([video[:8], video[-8:]]) generated_audio_slice = torch.cat([audio[:8], audio[-8:]]) - assert torch.allclose(expected_video_slice, generated_video_slice, atol=1e-4, rtol=1e-4) - assert torch.allclose(expected_audio_slice, generated_audio_slice, atol=1e-4, rtol=1e-4) + assert_tensors_close(generated_video_slice, expected_video_slice, atol=1e-4, rtol=1e-4) + assert_tensors_close(generated_audio_slice, expected_audio_slice, atol=1e-4, rtol=1e-4) def test_two_stages_inference(self): - device = "cpu" + # Run on CPU: the expected slices below are CPU-specific. + pipe = self.get_pipeline() - 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) + inputs = self.get_dummy_inputs() inputs["output_type"] = "latent" first_stage_output = pipe(**inputs) video_latent = first_stage_output.frames audio_latent = first_stage_output.audio - self.assertEqual(video_latent.shape, (1, 4, 3, 16, 16)) - self.assertEqual(audio_latent.shape, (1, 2, 5, 2)) - self.assertEqual(audio_latent.shape[1], components["vocoder"].config.out_channels) + assert video_latent.shape == (1, 4, 3, 16, 16) + assert audio_latent.shape == (1, 2, 5, 2) + assert audio_latent.shape[1] == pipe.vocoder.config.out_channels inputs["latents"] = video_latent inputs["audio_latents"] = audio_latent @@ -273,9 +280,9 @@ def test_two_stages_inference(self): video = second_stage_output.frames audio = second_stage_output.audio - self.assertEqual(video.shape, (1, 5, 3, 32, 32)) - self.assertEqual(audio.shape[0], 1) - self.assertEqual(audio.shape[1], components["vocoder"].config.out_channels) + assert video.shape == (1, *self.output_shape) + assert audio.shape[0] == 1 + assert audio.shape[1] == pipe.vocoder.config.out_channels # fmt: off expected_video_slice = torch.tensor( @@ -295,31 +302,16 @@ def test_two_stages_inference(self): generated_video_slice = torch.cat([video[:8], video[-8:]]) generated_audio_slice = torch.cat([audio[:8], audio[-8:]]) - assert torch.allclose(expected_video_slice, generated_video_slice, atol=1e-4, rtol=1e-4) - assert torch.allclose(expected_audio_slice, generated_audio_slice, atol=1e-4, rtol=1e-4) + assert_tensors_close(generated_video_slice, expected_video_slice, atol=1e-4, rtol=1e-4) + assert_tensors_close(generated_audio_slice, expected_audio_slice, atol=1e-4, rtol=1e-4) - def test_inference_batch_single_identical(self): - self._test_inference_batch_single_identical(batch_size=2, expected_max_diff=2e-2) - - def get_dummy_duration_head(self): - torch.manual_seed(0) - # The dummy connectors emit 4 heads * 8 head_dim = 32 wide output for both streams. - return LTX2DurationHead( - video_cross_attention_dim=32, - audio_cross_attention_dim=32, - pooler_hidden_dim=8, - num_queries=1, - num_pooler_heads=2, - mlp_hidden_dim=8, - ) + def test_inference_batch_single_identical(self, batch_size=2, expected_max_diff=2e-2): + super().test_inference_batch_single_identical(batch_size=batch_size, expected_max_diff=expected_max_diff) def test_auto_duration_produces_a_grid_valid_frame_count(self): - components = self.get_dummy_components() - components["duration_head"] = self.get_dummy_duration_head() - pipe = self.pipeline_class(**components).to(torch_device) - pipe.set_progress_bar_config(disable=None) + pipe = self.get_pipeline_with_duration_head() - inputs = self.get_dummy_inputs(torch_device) + inputs = self.get_dummy_inputs() inputs.pop("num_frames") inputs["min_seconds"] = 0.5 inputs["max_seconds"] = 2.0 @@ -330,19 +322,16 @@ def test_auto_duration_produces_a_grid_valid_frame_count(self): assert 0 < len(frames) <= round(2.0 * inputs["frame_rate"]) def test_omitting_num_frames_auto_predicts_when_a_head_is_present(self): - components = self.get_dummy_components() - components["duration_head"] = self.get_dummy_duration_head() - pipe = self.pipeline_class(**components).to(torch_device) - pipe.set_progress_bar_config(disable=None) + pipe = self.get_pipeline_with_duration_head() - inputs = self.get_dummy_inputs(torch_device) + inputs = self.get_dummy_inputs() inputs.pop("num_frames") inputs["min_seconds"] = 0.5 inputs["max_seconds"] = 2.0 bounded_frames = pipe(**inputs).frames[0] # Omitting `num_frames` entirely with default bounds must also take the auto path. - inputs = self.get_dummy_inputs(torch_device) + inputs = self.get_dummy_inputs() inputs.pop("num_frames") default_frames = pipe(**inputs).frames[0] @@ -355,10 +344,9 @@ def test_omitting_num_frames_uses_the_legacy_default_without_a_head(self): # Guards backwards compatibility: a pre-2.5 pipeline has no duration_head and must keep 121. components = self.get_dummy_components() assert components.get("duration_head") is None - pipe = self.pipeline_class(**components).to(torch_device) - pipe.set_progress_bar_config(disable=None) + pipe = self.get_pipeline(**components).to(torch_device) - inputs = self.get_dummy_inputs(torch_device) + inputs = self.get_dummy_inputs() inputs.pop("num_frames") # Decoding 121 frames is needlessly slow here; the latent frame count already pins num_frames down. inputs["output_type"] = "latent" @@ -369,12 +357,9 @@ def test_omitting_num_frames_uses_the_legacy_default_without_a_head(self): assert latents.shape[2] == expected_latent_frames def test_explicit_num_frames_wins_over_a_present_head(self): - components = self.get_dummy_components() - components["duration_head"] = self.get_dummy_duration_head() - pipe = self.pipeline_class(**components).to(torch_device) - pipe.set_progress_bar_config(disable=None) + pipe = self.get_pipeline_with_duration_head() - inputs = self.get_dummy_inputs(torch_device) + inputs = self.get_dummy_inputs() inputs["num_frames"] = 9 frames = pipe(**inputs).frames[0] @@ -383,28 +368,21 @@ def test_explicit_num_frames_wins_over_a_present_head(self): def test_auto_duration_with_multiple_prompts_raises(self): # The head predicts one duration, so it cannot serve prompts with different natural lengths. # Without this guard the pipeline silently applied the first prompt's length to all of them. - components = self.get_dummy_components() - components["duration_head"] = self.get_dummy_duration_head() - pipe = self.pipeline_class(**components).to(torch_device) - pipe.set_progress_bar_config(disable=None) + pipe = self.get_pipeline_with_duration_head() - inputs = self.get_dummy_inputs(torch_device) + inputs = self.get_dummy_inputs() inputs["prompt"] = ["a robot dancing", "a much longer and quite different scene"] inputs["negative_prompt"] = ["", ""] inputs.pop("num_frames") - with self.assertRaises(ValueError) as ctx: + with pytest.raises(ValueError, match="2 prompts were supplied"): pipe(**inputs) - assert "2 prompts were supplied" in str(ctx.exception) def test_multiple_prompts_still_work_with_an_explicit_num_frames(self): # The guard must be scoped to the auto path -- batched prompts with an integer are unaffected. - components = self.get_dummy_components() - components["duration_head"] = self.get_dummy_duration_head() - pipe = self.pipeline_class(**components).to(torch_device) - pipe.set_progress_bar_config(disable=None) + pipe = self.get_pipeline_with_duration_head() - inputs = self.get_dummy_inputs(torch_device) + inputs = self.get_dummy_inputs() inputs["prompt"] = ["a robot dancing", "a much longer and quite different scene"] inputs["negative_prompt"] = ["", ""] inputs["num_frames"] = 5 @@ -415,16 +393,40 @@ def test_multiple_prompts_still_work_with_an_explicit_num_frames(self): assert latents.shape[0] == 2 def test_invalid_duration_bounds_raise(self): - components = self.get_dummy_components() - components["duration_head"] = self.get_dummy_duration_head() - pipe = self.pipeline_class(**components).to(torch_device) - pipe.set_progress_bar_config(disable=None) + pipe = self.get_pipeline_with_duration_head() - inputs = self.get_dummy_inputs(torch_device) + inputs = self.get_dummy_inputs() inputs.pop("num_frames") inputs["min_seconds"] = 5.0 inputs["max_seconds"] = 2.0 - with self.assertRaises(ValueError) as ctx: + with pytest.raises(ValueError, match="min_seconds"): pipe(**inputs) - assert "min_seconds" in str(ctx.exception) + + +class TestLTX2PipelineMemory(LTX2PipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the LTX2 pipeline.""" + + @require_torch_accelerator + def test_group_offloading_inference(self): + # The shared helper only offloads a fixed set of component names and leaves LTX2's extra module + # components (`connectors`, `audio_vae`, `vocoder`) on CPU, so the forward pass mixes devices. + # Pipeline-level offloading, which walks every component, is exercised by + # `test_pipeline_level_group_offloading_inference`. + pytest.skip("Using test_pipeline_level_group_offloading_inference instead") + + +class TestLTX2PipelineLoRA(LTX2PipelineTesterConfig, LoraTesterMixin): + """LoRA tests for the LTX2 pipeline.""" + + # `LTX2Pipeline` advertises `connectors` as LoRA-loadable and `load_lora_weights` does handle connector + # LoRAs, but `save_lora_weights` only accepts `transformer_lora_layers` — so connector adapters cannot + # round-trip through the public API these tests drive. Scope the tests to the transformer until that closes. + lora_loadable_components = ["transformer"] + + +class TestLTX2PipelineLoRAMemory(LTX2PipelineTesterConfig, LoraMemoryTesterMixin): + """LoRA x memory-optimization tests (group offload, CPU offload) for the LTX2 pipeline.""" + + # See `TestLTX2PipelineLoRA`. + lora_loadable_components = ["transformer"] diff --git a/tests/pipelines/testing_utils/lora.py b/tests/pipelines/testing_utils/lora.py index 15af20a9abdb..4fbdfd066f85 100644 --- a/tests/pipelines/testing_utils/lora.py +++ b/tests/pipelines/testing_utils/lora.py @@ -106,15 +106,25 @@ def setup_method(self): if not issubclass(self.pipeline_class, LoraBaseMixin): pytest.skip(f"LoRA is not supported for this pipeline ({self.pipeline_class.__name__}).") + @property + def lora_loadable_components(self): + """Pipeline components these tests attach adapters to. + + Defaults to everything the pipeline advertises as LoRA-loadable. Override with a plain list on a test + class to narrow it — e.g. when a component is loadable but `save_lora_weights` cannot round-trip it yet, + so the save/load tests here could never pass for it. + """ + return self.pipeline_class._lora_loadable_modules + @property def text_encoder_components(self): """Names of the pipeline's LoRA-loadable text encoders, e.g. `["text_encoder", "text_encoder_2"]`.""" - return [name for name in self.pipeline_class._lora_loadable_modules if name.startswith("text_encoder")] + return [name for name in self.lora_loadable_components if name.startswith("text_encoder")] @property def denoiser_components(self): """Names of the pipeline's LoRA-loadable denoisers, e.g. `["unet"]` or `["transformer"]`.""" - return [name for name in self.pipeline_class._lora_loadable_modules if not name.startswith("text_encoder")] + return [name for name in self.lora_loadable_components if not name.startswith("text_encoder")] def get_denoiser(self, pipe): return pipe.transformer if hasattr(pipe, "transformer") else pipe.unet @@ -144,7 +154,7 @@ def add_adapters_to_pipeline(self, pipe, components=None, adapter_name="default" Returns {component_name: module} for everything adapted, e.g. for passing to `save_lora_weights` via `_get_lora_state_dicts`. """ - components = components if components is not None else self.pipeline_class._lora_loadable_modules + components = components if components is not None else self.lora_loadable_components adapted = {} for name in components: module = getattr(pipe, name, None)