Fix graph breaks

Author: Chillee

src/diffusers/models/attention.py

@@ -192,7 +192,9 @@ def forward(self, hidden_states):
         hidden_states = (hidden_states + residual) / self.rescale_output_factor
         return hidden_states
 
+from torch._dynamo import allow_in_graph
 
+@allow_in_graph
 class BasicTransformerBlock(nn.Module):
     r"""
     A basic Transformer block.

src/diffusers/models/modeling_utils.py

@@ -77,6 +77,7 @@ def find_tensor_attributes(module: torch.nn.Module) -> List[Tuple[str, Tensor]]:
 
 def get_parameter_dtype(parameter: torch.nn.Module):
     try:
+        return tuple(parameter.parameters())[0].dtype
         parameters_and_buffers = itertools.chain(parameter.parameters(), parameter.buffers())
         return next(parameters_and_buffers).dtype
     except StopIteration:

src/diffusers/models/unet_2d_blocks.py

@@ -560,7 +560,8 @@ def forward(
                 hidden_states,
                 encoder_hidden_states=encoder_hidden_states,
                 cross_attention_kwargs=cross_attention_kwargs,
-            ).sample
+                return_dict=False,
+            )[0]
             hidden_states = resnet(hidden_states, temb)
 
         return hidden_states
@@ -868,7 +869,8 @@ def custom_forward(*inputs):
                     hidden_states,
                     encoder_hidden_states=encoder_hidden_states,
                     cross_attention_kwargs=cross_attention_kwargs,
-                ).sample
+                    return_dict=False
+                )[0]
 
             output_states += (hidden_states,)
 
@@ -1859,7 +1861,8 @@ def custom_forward(*inputs):
                     hidden_states,
                     encoder_hidden_states=encoder_hidden_states,
                     cross_attention_kwargs=cross_attention_kwargs,
-                ).sample
+                    return_dict=False
+                )[0]
 
         if self.upsamplers is not None:
             for upsampler in self.upsamplers:

src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py

@@ -440,7 +440,7 @@ def run_safety_checker(self, image, device, dtype):
 
     def decode_latents(self, latents):
         latents = 1 / self.vae.config.scaling_factor * latents
-        image = self.vae.decode(latents).sample
+        image = self.vae.decode(latents, return_dict=False)[0]
         image = (image / 2 + 0.5).clamp(0, 1)
         # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
         image = image.cpu().permute(0, 2, 3, 1).float().numpy()
@@ -686,15 +686,16 @@ def __call__(
                     t,
                     encoder_hidden_states=prompt_embeds,
                     cross_attention_kwargs=cross_attention_kwargs,
-                ).sample
+                    return_dict=False,
+                )[0]
 
                 # perform guidance
                 if do_classifier_free_guidance:
                     noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
                     noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
 
                 # compute the previous noisy sample x_t -> x_t-1
-                latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample
+                latents = self.scheduler.step(noise_pred, t, latents, return_dict=False, **extra_step_kwargs)[0]
 
                 # call the callback, if provided
                 if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):