Update app.py

app.py

@@ -5,6 +5,8 @@ import tempfile
 import subprocess
 from pathlib import Path
 
+from dataclasses import dataclass
+import torch.multiprocessing as mp
 import spaces
 import gradio as gr
 import torch
@@ -205,15 +207,27 @@ def create_demo(
 
     return demo
 
-if __name__ == "__main__":
-    import argparse
+@dataclass
+class Config:
+    name: str = "flux-dev"
+    device: str = "cuda" if torch.cuda.is_available() else "cpu"
+    offload: bool = False
+    share: bool = False
+    ckpt_dir: str = "."
+    
+def parse_args() -> Config:
     parser = argparse.ArgumentParser(description="Flux")
     parser.add_argument("--name", type=str, default="flux-dev", help="Model name")
     parser.add_argument("--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu", help="Device to use")
     parser.add_argument("--offload", action="store_true", help="Offload model to CPU when not in use")
     parser.add_argument("--share", action="store_true", help="Create a public link to your demo")
     parser.add_argument("--ckpt_dir", type=str, default=".", help="Folder with checkpoints in safetensors format")
+    
     args = parser.parse_args()
+    return Config(**vars(args))
 
-    demo = create_demo(args.name, args.device, args.offload, args.ckpt_dir)
-    demo.launch(share=args.share)
+if __name__ == "__main__":
+    mp.set_start_method("spawn")
+    config = Config()
+    demo = create_demo(config.name, config.device, config.offload, config.ckpt_dir)
+    demo.launch(share=config.share)