Update src/flux/xflux_pipeline.py

src/flux/xflux_pipeline.py

@@ -1,3 +1,5 @@
+import torch.multiprocessing as mp
+
 from PIL import Image, ExifTags
 import numpy as np
 import torch
@@ -34,6 +36,7 @@ import spaces
 
 class XFluxPipeline:
     def __init__(self, model_type, device, offload: bool = False):
+        mp.set_start_method('spawn', force=True)
         self.device = torch.device(device)
         self.offload = offload
         self.model_type = model_type
@@ -162,7 +165,6 @@ class XFluxPipeline:
         image_proj = self.improj(image_prompt_embeds)
         return image_proj
 
-    @spaces.GPU()
     def __call__(self,
                  prompt: str,
                  image_prompt: Image = None,
@@ -219,6 +221,7 @@ 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,