Upload matryoshka.py

matryoshka.py

@@ -3762,8 +3762,7 @@ class MatryoshkaPipeline(
     """
 
     model_cpu_offload_seq = "text_encoder->image_encoder->unet"
-    _optional_components = ["unet", "safety_checker", "feature_extractor", "image_encoder"]
-    _exclude_from_cpu_offload = ["safety_checker"]
+    _optional_components = ["unet", "feature_extractor", "image_encoder"]
     _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"]
 
     def __init__(
@@ -4549,18 +4548,10 @@ class MatryoshkaPipeline(
                 self.do_classifier_free_guidance,
             )
 
-        # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
-        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
         # 4. Prepare timesteps
-        if isinstance(self.scheduler, MatryoshkaDDIMScheduler):
-            timesteps, num_inference_steps = retrieve_timesteps(
-                self.scheduler, num_inference_steps, device, timesteps, sigmas
-            )
-
-            extra_step_kwargs |= {"use_clipped_model_output": True}
-        else:
-            timesteps = self.scheduler.timesteps
-
+        timesteps, num_inference_steps = retrieve_timesteps(
+            self.scheduler, num_inference_steps, device, timesteps, sigmas
+        )
         timesteps = timesteps[:-1]
 
         # 5. Prepare latent variables
@@ -4577,6 +4568,10 @@ class MatryoshkaPipeline(
             latents,
         )
 
+        # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
+        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
+        extra_step_kwargs |= {"use_clipped_model_output": True}
+
         # 6.1 Add image embeds for IP-Adapter
         added_cond_kwargs = (
             {"image_embeds": image_embeds}
@@ -4634,19 +4629,7 @@ class MatryoshkaPipeline(
                     noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=self.guidance_rescale)
 
                 # compute the previous noisy sample x_t -> x_t-1
-                if self.scheduler.scales is not None and not isinstance(self.scheduler, MatryoshkaDDIMScheduler):
-                    latents[0] = self.scheduler.step(
-                        noise_pred[0], t, latents[0], **extra_step_kwargs, return_dict=False
-                    )[0]
-                    latents[1] = self.scheduler.inner_scheduler.step(
-                        noise_pred[1], t, latents[1], **extra_step_kwargs, return_dict=False
-                    )[0]
-                    if len(latents) > 2:
-                        latents[2] = self.scheduler.inner_scheduler.inner_scheduler.step(
-                            noise_pred[2], t, latents[2], **extra_step_kwargs, return_dict=False
-                        )[0]
-                else:
-                    latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
+                latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
 
                 if callback_on_step_end is not None:
                     callback_kwargs = {}
@@ -4672,8 +4655,7 @@ class MatryoshkaPipeline(
 
         if self.scheduler.scales is not None:
             for i, (img, scale) in enumerate(zip(image, self.scheduler.scales)):
-                img = torch.clip(img * scale, -1, 1)
-                image[i] = self.image_processor.postprocess(img, output_type=output_type)[0]
+                image[i] = self.image_processor.postprocess(img * scale, output_type=output_type)[0]
         else:
             image = self.image_processor.postprocess(image, output_type=output_type)