Update custom_pipeline.py

custom_pipeline.py

@@ -47,10 +47,6 @@ class FluxWithCFGPipeline(FluxPipeline):
     Extends the FluxPipeline to yield intermediate images during the denoising process 
     with progressively increasing resolution for faster generation.
     """
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.default_sample_size = 512  # Default sample size from the first pipeline
-
     @torch.inference_mode()
     def generate_images(
         self,
@@ -106,7 +102,6 @@ class FluxWithCFGPipeline(FluxPipeline):
             max_sequence_length=max_sequence_length,
             lora_scale=lora_scale,
         )
-
         # 4. Prepare latent variables
         num_channels_latents = self.transformer.config.in_channels // 4
         latents, latent_image_ids = self.prepare_latents(
@@ -119,7 +114,6 @@ class FluxWithCFGPipeline(FluxPipeline):
             generator,
             latents,
         )
-
         # 5. Prepare timesteps
         sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps)
         image_seq_len = latents.shape[1]
@@ -156,12 +150,14 @@ class FluxWithCFGPipeline(FluxPipeline):
                 return_dict=False,
             )[0]
 
-            # Yield intermediate result
+             # Yield intermediate result
             latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
             torch.cuda.empty_cache()
 
         # Final image
         return self._decode_latents_to_image(latents, height, width, output_type)
+        self.maybe_free_model_hooks()
+        torch.cuda.empty_cache()
 
     def _decode_latents_to_image(self, latents, height, width, output_type, vae=None):
         """Decodes the given latents into an image."""