Upload app.py with huggingface_hub

app.py

@@ -6,26 +6,41 @@ Interactive web UI for image generation
 import gradio as gr
 from bytedream.generator import ByteDreamGenerator
 import torch
+import os
 
 
 # Initialize generator
 print("Loading Byte Dream model...")
+
+# Check if we should load from Hugging Face
+HF_REPO_ID = os.getenv("HF_REPO_ID", None)
+MODEL_PATH = os.getenv("MODEL_PATH", "./models/bytedream")
+
 try:
-    generator = ByteDreamGenerator(
-        model_path="./models/bytedream",
-        config_path="config.yaml",
-        device="cpu",
-    )
+    if HF_REPO_ID:
+        print(f"Loading model from Hugging Face: {HF_REPO_ID}...")
+        generator = ByteDreamGenerator(
+            hf_repo_id=HF_REPO_ID,
+            config_path="config.yaml",
+            device="cpu",
+        )
+    else:
+        print(f"Loading model from local path: {MODEL_PATH}...")
+        generator = ByteDreamGenerator(
+            model_path=MODEL_PATH,
+            config_path="config.yaml",
+            device="cpu",
+        )
     print("✓ Model loaded successfully!")
 except Exception as e:
     print(f"⚠ Warning: Could not load model: {e}")
     print("  Please train the model first using: python train.py")
     print("  Or download pretrained weights from Hugging Face.")
     print("")
-    print("  To use a model from Hugging Face, run:")
-    print("    python infer.py --prompt 'your prompt' --model 'username/repo_name'")
+    print("  To use a model from Hugging Face, set environment variable:")
+    print("    HF_REPO_ID=username/repo_name")
     print("")
-    print("Starting in demo mode with random initialization...")
+    print("Starting in demo mode...")
     generator = None