[Bugfix] Fix Dtype Crashes in SD3 (#2526)

Signed-off-by: Alex Brooks <albrooks@redhat.com>
Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>

vllm_omni/diffusion/models/sd3/pipeline_sd3.py

@@ -321,10 +321,8 @@ class StableDiffusion3Pipeline(nn.Module, CFGParallelMixin, DiffusionPipelinePro
         self,
         prompt: str | list[str] = "",
         num_images_per_prompt: int = 1,
-        dtype: torch.dtype | None = None,
         clip_model_index: int = 0,
     ):
-        dtype = dtype or self.text_encoder.dtype
         clip_tokenizers = [self.tokenizer, self.tokenizer_2]
         clip_text_encoders = [self.text_encoder, self.text_encoder_2]
 
@@ -351,11 +349,11 @@ class StableDiffusion3Pipeline(nn.Module, CFGParallelMixin, DiffusionPipelinePro
                 f" {self.tokenizer_max_length} tokens: {removed_text}"
             )
         prompt_embeds = text_encoder(text_input_ids.to(self.device), output_hidden_states=True)
-        pooled_prompt_embeds = prompt_embeds[0]
+        pooled_prompt_embeds = prompt_embeds[0].to(dtype=self.od_config.dtype, device=self.device)
 
         prompt_embeds = prompt_embeds.hidden_states[-2]
 
-        prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=self.device)
+        prompt_embeds = prompt_embeds.to(dtype=self.od_config.dtype, device=self.device)
         _, seq_len, _ = prompt_embeds.shape
         prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
         prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
@@ -369,10 +367,7 @@ class StableDiffusion3Pipeline(nn.Module, CFGParallelMixin, DiffusionPipelinePro
         prompt: str | list[str] = "",
         num_images_per_prompt: int = 1,
         max_sequence_length: int = 256,
-        dtype: torch.dtype | None = None,
-    ):
-        dtype = dtype or self.text_encoder_3.dtype
-
+    ) -> torch.Tensor:
         prompt = [prompt] if isinstance(prompt, str) else prompt
         batch_size = len(prompt)
 
@@ -384,7 +379,7 @@ class StableDiffusion3Pipeline(nn.Module, CFGParallelMixin, DiffusionPipelinePro
                     self.transformer.joint_attention_dim,
                 ),
                 device=self.device,
-                dtype=dtype,
+                dtype=self.od_config.dtype,
             )
 
         text_inputs = self.tokenizer_3(
@@ -407,8 +402,7 @@ class StableDiffusion3Pipeline(nn.Module, CFGParallelMixin, DiffusionPipelinePro
 
         prompt_embeds = self.text_encoder_3(text_input_ids.to(self.device))[0]
 
-        dtype = self.text_encoder_3.dtype
-        prompt_embeds = prompt_embeds.to(dtype=dtype, device=self.device)
+        prompt_embeds = prompt_embeds.to(dtype=self.od_config.dtype, device=self.device)
         _, seq_len, _ = prompt_embeds.shape
         prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
         prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
@@ -485,13 +479,11 @@ class StableDiffusion3Pipeline(nn.Module, CFGParallelMixin, DiffusionPipelinePro
         num_channels_latents,
         height,
         width,
-        dtype,
-        device,
         generator,
         latents=None,
     ) -> torch.Tensor:
         if latents is not None:
-            return latents.to(device=device, dtype=dtype)
+            return latents.to(device=self.device, dtype=self.od_config.dtype)
 
         shape = (
             batch_size,
@@ -506,7 +498,7 @@ class StableDiffusion3Pipeline(nn.Module, CFGParallelMixin, DiffusionPipelinePro
                 f" size of {batch_size}. Make sure the batch size matches the length of the generators."
             )
 
-        latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+        latents = randn_tensor(shape, generator=generator, device=self.device, dtype=self.od_config.dtype)
 
         return latents
 
@@ -550,9 +542,9 @@ class StableDiffusion3Pipeline(nn.Module, CFGParallelMixin, DiffusionPipelinePro
         latents: torch.Tensor,
         timesteps: torch.Tensor,
         prompt_embeds: torch.Tensor,
-        pooled_prompt_embeds: torch.Tensor,
-        negative_prompt_embeds: torch.Tensor,
-        negative_pooled_prompt_embeds: torch.Tensor,
+        pooled_prompt_embeds: torch.Tensor | None,
+        negative_prompt_embeds: torch.Tensor | None,
+        negative_pooled_prompt_embeds: torch.Tensor | None,
         do_true_cfg: bool,
         guidance_scale: float,
         cfg_normalize: bool = False,
@@ -582,8 +574,7 @@ class StableDiffusion3Pipeline(nn.Module, CFGParallelMixin, DiffusionPipelinePro
             self._current_timestep = t
 
             # Broadcast timestep to match batch size
-            timestep = t.expand(latents.shape[0]).to(device=latents.device, dtype=latents.dtype)
-
+            timestep = t.expand(latents.shape[0]).to(device=self.device, dtype=self.od_config.dtype)
             positive_kwargs = {
                 "hidden_states": latents,
                 "timestep": timestep,
@@ -711,8 +702,6 @@ class StableDiffusion3Pipeline(nn.Module, CFGParallelMixin, DiffusionPipelinePro
             num_channels_latents,
             height,
             width,
-            prompt_embeds.dtype,
-            self.device,
             generator,
             latents,
         )
@@ -737,6 +726,7 @@ class StableDiffusion3Pipeline(nn.Module, CFGParallelMixin, DiffusionPipelinePro
         if self.output_type == "latent":
             image = latents
         else:
+            # Ensure the latents are the same dtype as the VAE for decode
             latents = latents.to(self.vae.dtype)
             latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor