Update live_preview_helpers.py

live_preview_helpers.py

@@ -2,7 +2,6 @@ import torch
 import numpy as np
 from diffusers import FluxPipeline, AutoencoderTiny, FlowMatchEulerDiscreteScheduler
 from typing import Any, Dict, List, Optional, Union
-import spaces
 
 # Helper functions
 def calculate_shift(
@@ -41,7 +40,7 @@ def retrieve_timesteps(
     return timesteps, num_inference_steps
 
 # FLUX pipeline function
-@spaces.GPU()
+@torch.inference_mode()
 def flux_pipe_call_that_returns_an_iterable_of_images(
     self,
     prompt: Union[str, List[str]] = None,
@@ -164,4 +163,4 @@ def flux_pipe_call_that_returns_an_iterable_of_images(
     image = good_vae.decode(latents, return_dict=False)[0]
     self.maybe_free_model_hooks()
     torch.cuda.empty_cache()
-    yield self.image_processor.postprocess(image, output_type=output_type)[0]
+    yield self.image_processor.postprocess(image, output_type=output_type)[0]