Update app.py

app.py

@@ -191,19 +191,36 @@ if __name__ == "__main__":
 
     @dataclasses.dataclass
     class AppArgs:
-        if local_model == True:
-            model_name: str = "ghibli-fine-tuned-sd-2.1"
-        else:
-            model_name: str = "danhtran2mind/ghibli-fine-tuned-sd-2.1"
-        device: str = "cuda" if torch.cuda.is_available() else "cpu"
-        port: int = 7860
-        share: bool = False  # Set to True for public sharing (Hugging Face Spaces)
+        local_model: bool = dataclasses.field(
+            default=False, metadata={"help": "Use local model path instead of Hugging Face model."}
+        )
+        model_name: str = dataclasses.field(
+            default="danhtran2mind/ghibli-fine-tuned-sd-2.1",
+            metadata={"help": "Model name or path for the fine-tuned Stable Diffusion model."}
+        )
+        device: str = dataclasses.field(
+            default="cuda" if torch.cuda.is_available() else "cpu",
+            metadata={"help": "Device to run the model on (e.g., 'cuda', 'cpu')."}
+        )
+        port: int = dataclasses.field(
+            default=7860, metadata={"help": "Port to run the Gradio app on."}
+        )
+        share: bool = dataclasses.field(
+            default=False, metadata={"help": "Set to True for public sharing (Hugging Face Spaces)."}
+        )
 
     parser = HfArgumentParser([AppArgs])
     args_tuple = parser.parse_args_into_dataclasses()
     args = args_tuple[0]
 
+    # Set model_name based on local_model flag
+    if args.local_model:
+        args.model_name = "ghibli-fine-tuned-sd-2.1"
+
+    demo = create_demo(args.model_name, args.device)
+    demo.launch(server_port=args.port, share=args.share)
+
     demo = create_demo(args.model_name, args.device)
     demo.launch(server_port=args.port, share=args.share)
 
-<<add option choose local_model when run app.py>>
+<<add option choose local_model when run app.py>>