Update custom_pipeline.py

custom_pipeline.py

@@ -130,60 +130,29 @@ class FluxWithCFGPipeline(FluxPipeline):
 
         # Handle guidance
         guidance = torch.full([1], guidance_scale, device=device, dtype=torch.float16).expand(latents.shape[0]) if self.transformer.config.guidance_embeds else None
-        
-        # static method that can be jitted
-        @staticmethod
-        @torch.jit.script
-        def _denoising_loop_static(latents, timesteps, pooled_prompt_embeds, prompt_embeds, text_ids, latent_image_ids, guidance, joint_attention_kwargs, transformer, scheduler):
-            for i, t in enumerate(timesteps):
-                timestep = t.expand(latents.shape[0]).to(latents.dtype)
-
-                noise_pred = transformer(
-                    hidden_states=latents,
-                    timestep=timestep / 1000,
-                    guidance=guidance,
-                    pooled_projections=pooled_prompt_embeds,
-                    encoder_hidden_states=prompt_embeds,
-                    txt_ids=text_ids,
-                    img_ids=latent_image_ids,
-                    joint_attention_kwargs=joint_attention_kwargs,
-                    return_dict=False,
-                )[0]
-
-                latents = scheduler.step(noise_pred, t, latents, return_dict=False)[0]
-                torch.cuda.empty_cache()
-            return latents
-
-        # Make the core denoising loop a static method
-        self._denoising_loop = torch.cuda.make_graphed_callables(
-            _denoising_loop_static,
-            (
-                latents.clone(),  # Example inputs for warmup
-                timesteps.clone(),
-                pooled_prompt_embeds.clone(),
-                prompt_embeds.clone(),
-                text_ids.clone(),
-                latent_image_ids.clone(),
-                guidance.clone(),
-                self._joint_attention_kwargs,
-                self.transformer,
-                self.scheduler
-            )
-        )
 
-        # Call the static method now
-        latents = self._denoising_loop(
-            latents,
-            timesteps,
-            pooled_prompt_embeds,
-            prompt_embeds,
-            text_ids,
-            latent_image_ids,
-            guidance,
-            self._joint_attention_kwargs,
-            self.transformer,
-            self.scheduler
-        )
+        # 6. Denoising loop
+        for i, t in enumerate(timesteps):
+            if self.interrupt:
+                continue
+
+            timestep = t.expand(latents.shape[0]).to(latents.dtype)
+
+            noise_pred = self.transformer(
+                hidden_states=latents,
+                timestep=timestep / 1000,
+                guidance=guidance,
+                pooled_projections=pooled_prompt_embeds,
+                encoder_hidden_states=prompt_embeds,
+                txt_ids=text_ids,
+                img_ids=latent_image_ids,
+                joint_attention_kwargs=self.joint_attention_kwargs,
+                return_dict=False,
+            )[0]
+
+             # 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)
@@ -196,4 +165,4 @@ class FluxWithCFGPipeline(FluxPipeline):
         latents = self._unpack_latents(latents, height, width, self.vae_scale_factor)
         latents = (latents / vae.config.scaling_factor) + vae.config.shift_factor
         image = vae.decode(latents, return_dict=False)[0]
-        return self.image_processor.postprocess(image, output_type=output_type)[0]
+        return self.image_processor.postprocess(image, output_type=output_type)[0]