Update app.py

app.py

@@ -1,20 +1,19 @@
-import spaces
+import os
 import torch
-
+import spaces
 import re
-import os
 import yaml
 import tempfile
 import subprocess
 from pathlib import Path
-
 from dataclasses import dataclass
-
 import gradio as gr
 from src.flux.xflux_pipeline import XFluxPipeline
-
-import os
 from huggingface_hub import login
+import multiprocessing
+
+# Set the multiprocessing start method globally
+multiprocessing.set_start_method('spawn', force=True)
 
 hf_token = os.getenv("HF_TOKEN")
 if hf_token:
@@ -29,14 +28,13 @@ class Config:
     offload: bool = False
     share: bool = False
     ckpt_dir: str = "."
-    
+
 xflux_pipeline = XFluxPipeline(Config.name, Config.device, Config.offload)
 xflux_pipeline.to(device='cuda' if torch.cuda.is_available() else 'cpu')
 
 @spaces.GPU
 @torch.inference_mode()
 def generate(**kwargs): 
-    torch.multiprocessing.set_start_method('spawn', force=True)
     return xflux_pipeline.gradio_generate(**kwargs)
     
 def parse_args() -> Config: