Update src/flux/xflux_pipeline.py

src/flux/xflux_pipeline.py

@@ -31,8 +31,6 @@ from src.flux.util import (
 )
 
 from transformers import CLIPVisionModelWithProjection, CLIPImageProcessor
-import spaces
-
 
 class XFluxPipeline:
     def __init__(self, model_type, device, offload: bool = False):
@@ -57,6 +55,12 @@ class XFluxPipeline:
         self.controlnet_loaded = False
         self.ip_loaded = False
 
+    def to(*args, **args):
+        self.model.to(*args, **kwargs)
+        self.clip.to(*args, **kwargs)
+        self.t5.to(*args, **kwargs)
+        self.ae.to(*args, **kwargs)
+        
     def set_ip(self, local_path: str = None, repo_id = None, name: str = None):
         self.model.to(self.device)
 
@@ -221,7 +225,6 @@ class XFluxPipeline:
             neg_ip_scale=neg_ip_scale,
         )
 
-    @spaces.GPU()
     def gradio_generate(self, prompt, image_prompt, controlnet_image, width, height, guidance,
                         num_steps, seed, true_gs, ip_scale, neg_ip_scale, neg_prompt,
                         neg_image_prompt, timestep_to_start_cfg, control_type, control_weight,