update

Author: DN6

src/diffusers/models/transformers/transformer_wan_vace.py

@@ -331,7 +331,7 @@ def forward(
                 )
                 if i in self.config.vace_layers:
                     control_hint, scale = control_hidden_states_list.pop()
-                    hidden_states = hidden_states + control_hint * scale
+                    hidden_states = hidden_states + control_hint.to(hidden_states.device) * scale
         else:
             # Prepare VACE hints
             control_hidden_states_list = []
@@ -346,7 +346,7 @@ def forward(
                 hidden_states = block(hidden_states, encoder_hidden_states, timestep_proj, rotary_emb)
                 if i in self.config.vace_layers:
                     control_hint, scale = control_hidden_states_list.pop()
-                    hidden_states = hidden_states + control_hint * scale
+                    hidden_states = hidden_states + control_hint.to(hidden_states.device) * scale
 
         # 6. Output norm, projection & unpatchify
         shift, scale = (self.scale_shift_table.to(temb.device) + temb.unsqueeze(1)).chunk(2, dim=1)

tests/models/transformers/test_models_transformer_wan.py

@@ -91,11 +91,13 @@ def get_dummy_inputs(self) -> dict[str, torch.Tensor]:
                 (batch_size, num_channels, num_frames, height, width),
                 generator=self.generator,
                 device=torch_device,
+                dtype=self.torch_dtype,
             ),
             "encoder_hidden_states": randn_tensor(
                 (batch_size, sequence_length, text_encoder_embedding_dim),
                 generator=self.generator,
                 device=torch_device,
+                dtype=self.torch_dtype,
             ),
             "timestep": torch.randint(0, 1000, size=(batch_size,), generator=self.generator).to(torch_device),
         }

tests/models/transformers/test_models_transformer_wan_animate.py

@@ -113,27 +113,32 @@ def get_dummy_inputs(self) -> dict[str, torch.Tensor]:
                 (batch_size, 2 * num_channels + 4, num_frames + 1, height, width),
                 generator=self.generator,
                 device=torch_device,
+                dtype=self.torch_dtype,
             ),
             "timestep": torch.randint(0, 1000, size=(batch_size,), generator=self.generator).to(torch_device),
             "encoder_hidden_states": randn_tensor(
                 (batch_size, sequence_length, text_encoder_embedding_dim),
                 generator=self.generator,
                 device=torch_device,
+                dtype=self.torch_dtype,
             ),
             "encoder_hidden_states_image": randn_tensor(
                 (batch_size, clip_seq_len, clip_dim),
                 generator=self.generator,
                 device=torch_device,
+                dtype=self.torch_dtype,
             ),
             "pose_hidden_states": randn_tensor(
                 (batch_size, num_channels, num_frames, height, width),
                 generator=self.generator,
                 device=torch_device,
+                dtype=self.torch_dtype,
             ),
             "face_pixel_values": randn_tensor(
                 (batch_size, 3, inference_segment_length, face_height, face_width),
                 generator=self.generator,
                 device=torch_device,
+                dtype=self.torch_dtype,
             ),
         }
 

tests/models/transformers/test_models_transformer_wan_vace.py

@@ -96,16 +96,19 @@ def get_dummy_inputs(self) -> dict[str, torch.Tensor]:
                 (batch_size, num_channels, num_frames, height, width),
                 generator=self.generator,
                 device=torch_device,
+                dtype=self.torch_dtype,
             ),
             "encoder_hidden_states": randn_tensor(
                 (batch_size, sequence_length, text_encoder_embedding_dim),
                 generator=self.generator,
                 device=torch_device,
+                dtype=self.torch_dtype,
             ),
             "control_hidden_states": randn_tensor(
                 (batch_size, vace_in_channels, num_frames, height, width),
                 generator=self.generator,
                 device=torch_device,
+                dtype=self.torch_dtype,
             ),
             "timestep": torch.randint(0, 1000, size=(batch_size,), generator=self.generator).to(torch_device),
         }