Update live_preview_helpers.py

live_preview_helpers.py

@@ -131,34 +131,31 @@ def flux_pipe_call_that_returns_an_iterable_of_images(
     guidance = torch.full([1], guidance_scale, device=device, dtype=torch.float32).expand(latents.shape[0]) if self.transformer.config.guidance_embeds else None
 
     # 6. Denoising loop
-    # FIX: The `para_attn` library requires wrapping the inference loop in a `cache_scope`
-    # context manager. This manager is attached to the pipeline (`self`) by `apply_cache_on_pipe`.
-    with self.cache_scope(0, num_inference_steps):
-        for i, t in enumerate(timesteps):
-            if self.interrupt:
-                continue
+    for i, t in enumerate(timesteps):
+        if self.interrupt:
+            continue
 
-            timestep = t.expand(latents.shape[0]).to(latents.dtype)
+        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_for_image = self._unpack_latents(latents, height, width, self.vae_scale_factor)
-            latents_for_image = (latents_for_image / self.vae.config.scaling_factor) + self.vae.config.shift_factor
-            image = self.vae.decode(latents_for_image, return_dict=False)[0]
-            yield self.image_processor.postprocess(image, output_type=output_type)[0]
-            
-            latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
-            torch.cuda.empty_cache()
+        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_for_image = self._unpack_latents(latents, height, width, self.vae_scale_factor)
+        latents_for_image = (latents_for_image / self.vae.config.scaling_factor) + self.vae.config.shift_factor
+        image = self.vae.decode(latents_for_image, return_dict=False)[0]
+        yield self.image_processor.postprocess(image, output_type=output_type)[0]
+        
+        latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
+        torch.cuda.empty_cache()
 
     # Final image using good_vae
     latents = self._unpack_latents(latents, height, width, self.vae_scale_factor)