Upload infer.py with huggingface_hub

infer.py

@@ -88,7 +88,14 @@ Examples:
         "--model", "-m",
         type=str,
         default=None,
-        help="Path to model directory (default: uses config)"
+        help="Path to model directory or Hugging Face repo ID (default: uses config)"
+    )
+    
+    parser.add_argument(
+        "--hf_repo",
+        type=str,
+        default=None,
+        help="Hugging Face repository ID to load model from (e.g., username/repo)"
     )
     
     parser.add_argument(
@@ -115,11 +122,20 @@ Examples:
     print("Byte Dream - AI Image Generator")
     print("="*60)
     
-    generator = ByteDreamGenerator(
-        model_path=args.model,
-        config_path=args.config,
-        device=args.device,
-    )
+    # Determine if loading from HF or local
+    if args.hf_repo:
+        print(f"Loading model from Hugging Face: {args.hf_repo}")
+        generator = ByteDreamGenerator(
+            hf_repo_id=args.hf_repo,
+            config_path=args.config,
+            device=args.device,
+        )
+    else:
+        generator = ByteDreamGenerator(
+            model_path=args.model,
+            config_path=args.config,
+            device=args.device,
+        )
     
     # Print model info
     info = generator.get_model_info()