Update src/flux/xflux_pipeline.py

src/flux/xflux_pipeline.py

@@ -29,7 +29,8 @@ from src.flux.util import (
 )
 
 from transformers import CLIPVisionModelWithProjection, CLIPImageProcessor
-
+import spaces
+`
 class XFluxPipeline:
     def __init__(self, model_type, device, offload: bool = False):
         self.device = torch.device(device)
@@ -160,6 +161,7 @@ class XFluxPipeline:
         image_proj = self.improj(image_prompt_embeds)
         return image_proj
 
+    @spaces.GPU
     def __call__(self,
                  prompt: str,
                  image_prompt: Image = None,