Update app.py

app.py

@@ -22,8 +22,28 @@ if hf_token:
     login(token=hf_token)
 else:
     print("No Hugging Face token found.")
+
+@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 = "."
     
+xflux_pipeline = XFluxPipeline(Config.name, Config.device, Config.offload)
 
+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))
+    
 def list_dirs(path):
     if path is None or path == "None" or path == "":
         return
@@ -124,7 +144,6 @@ def create_demo(
         offload: bool = False,
         ckpt_dir: str = "",
     ):
-    xflux_pipeline = XFluxPipeline(model_type, device, offload)
     checkpoints = sorted(Path(ckpt_dir).glob("*.safetensors"))
 
     with gr.Blocks() as demo:
@@ -208,25 +227,6 @@ def create_demo(
 
     return demo
 
-@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))
-
 if __name__ == "__main__":
     import torch.multiprocessing as mp
     mp.set_start_method('spawn', force=True)  # Corrected start method for CUDA