| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
|
|
| import unittest |
|
|
| import numpy as np |
| import torch |
| from transformers import AutoTokenizer, T5EncoderModel |
|
|
| from diffusers import ( |
| AutoencoderKLWan, |
| SkyReelsV2Pipeline, |
| SkyReelsV2Transformer3DModel, |
| UniPCMultistepScheduler, |
| ) |
|
|
| from ...testing_utils import ( |
| enable_full_determinism, |
| ) |
| from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS |
| from ..test_pipelines_common import ( |
| PipelineTesterMixin, |
| ) |
|
|
|
|
| enable_full_determinism() |
|
|
|
|
| class SkyReelsV2PipelineFastTests(PipelineTesterMixin, unittest.TestCase): |
| pipeline_class = SkyReelsV2Pipeline |
| 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", |
| ] |
| ) |
| test_xformers_attention = False |
| supports_dduf = False |
|
|
| def get_dummy_components(self): |
| 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], |
| ) |
|
|
| torch.manual_seed(0) |
| scheduler = UniPCMultistepScheduler(flow_shift=8.0, use_flow_sigmas=True) |
| text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5") |
| tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") |
|
|
| torch.manual_seed(0) |
| transformer = SkyReelsV2Transformer3DModel( |
| patch_size=(1, 2, 2), |
| num_attention_heads=2, |
| attention_head_dim=12, |
| in_channels=16, |
| out_channels=16, |
| text_dim=32, |
| freq_dim=256, |
| ffn_dim=32, |
| num_layers=2, |
| cross_attn_norm=True, |
| qk_norm="rms_norm_across_heads", |
| rope_max_seq_len=32, |
| ) |
|
|
| components = { |
| "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 = { |
| "prompt": "dance monkey", |
| "negative_prompt": "negative", |
| "generator": generator, |
| "num_inference_steps": 2, |
| "guidance_scale": 6.0, |
| "height": 16, |
| "width": 16, |
| "num_frames": 9, |
| "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) |
| video = pipe(**inputs).frames |
| generated_video = video[0] |
|
|
| self.assertEqual(generated_video.shape, (9, 3, 16, 16)) |
| expected_video = torch.randn(9, 3, 16, 16) |
| max_diff = np.abs(generated_video - expected_video).max() |
| self.assertLessEqual(max_diff, 1e10) |
|
|
| @unittest.skip("Test not supported") |
| def test_attention_slicing_forward_pass(self): |
| pass |
|
|
