Update live_preview_helpers.py

live_preview_helpers.py

@@ -60,7 +60,7 @@ 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, # *** FIX 3: Accept the context scope as a parameter ***
+    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,9 +132,8 @@ 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 4: Define a generator and use the passed-in cache_scope with it ***
     def denoise_loop_generator():
-        nonlocal latents # Allow this generator to modify the 'latents' variable from the parent scope
+        nonlocal latents
         for i, t in enumerate(timesteps):
             if self.interrupt:
                 continue
@@ -161,11 +160,10 @@ def flux_pipe_call_that_returns_an_iterable_of_images(
             latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
             torch.cuda.empty_cache()
 
-    # Use the context manager if it was provided
     if cache_scope:
-        with cache_scope(0, num_inference_steps):
+        with cache_scope:
             yield from denoise_loop_generator()
-    else: # Fallback if cache is not used
+    else:
         yield from denoise_loop_generator()
 
     # Final image using good_vae