Update live_preview_helpers.py

live_preview_helpers.py

@@ -60,7 +60,6 @@ def flux_pipe_call_that_returns_an_iterable_of_images(
     joint_attention_kwargs: Optional[Dict[str, Any]] = None,
     max_sequence_length: int = 512,
     good_vae: Optional[Any] = None,
-    cache_scope: Optional[Any] = None,
 ):
     height = height or self.default_sample_size * self.vae_scale_factor
     width = width or self.default_sample_size * self.vae_scale_factor
@@ -132,39 +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
-    def denoise_loop_generator():
-        nonlocal latents
-        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()
-
-    if cache_scope:
-        with cache_scope:
-            yield from denoise_loop_generator()
-    else:
-        yield from denoise_loop_generator()
+        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)