Upload quick_start.py with huggingface_hub

quick_start.py

@@ -1,124 +1,222 @@
 """
-Quick Start Script
-Setup and test Byte Dream in one command
+Quick Start Script for Hugging Face Integration
+Helps you upload/download models easily
 """
 
-import subprocess
 import sys
 from pathlib import Path
 
 
-def check_requirements():
-    """Check if requirements are installed"""
-    print("Checking requirements...")
-    
-    required = [
-        'torch',
-        'transformers',
-        'diffusers',
-        'pillow',
-        'numpy',
-        'gradio',
-    ]
-    
-    missing = []
-    
-    for package in required:
-        try:
-            __import__(package.replace('-', '_'))
-            print(f"  ✓ {package}")
-        except ImportError:
-            print(f"  ✗ {package} - MISSING")
-            missing.append(package)
-    
-    if missing:
-        print(f"\nMissing packages: {', '.join(missing)}")
-        print("\nInstall with:")
-        print("  pip install -r requirements.txt")
+def print_banner(text):
+    """Print formatted banner"""
+    print("\n" + "="*60)
+    print(text.center(60))
+    print("="*60 + "\n")
+
+
+def check_model_exists():
+    """Check if trained model exists"""
+    model_path = Path("./models/bytedream")
+    
+    if not model_path.exists():
+        print("❌ Model directory not found!")
+        print("\nPlease train the model first:")
+        print("  python train.py")
+        print("\nOr download from Hugging Face:")
+        print("  python infer.py --hf_repo username/repo --prompt 'test'")
+        return False
+    
+    # Check for weights
+    unet_weights = model_path / "unet" / "pytorch_model.bin"
+    vae_weights = model_path / "vae" / "pytorch_model.bin"
+    
+    if not (unet_weights.exists() or (model_path / "pytorch_model.bin").exists()):
+        print("⚠ Model directory exists but no weights found!")
+        print("Please train the model first.")
         return False
     
-    print("\n✓ All requirements satisfied!")
     return True
 
 
-def test_model():
-    """Test model generation"""
-    print("\n" + "="*60)
-    print("Testing Byte Dream Model")
-    print("="*60)
+def upload_to_hf():
+    """Upload model to Hugging Face"""
+    print_banner("UPLOAD TO HUGGING FACE HUB")
+    
+    # Check model exists
+    if not check_model_exists():
+        return
+    
+    # Get token
+    token = input("Enter your Hugging Face token (hf_...): ").strip()
+    if not token:
+        print("❌ Token is required!")
+        return
+    
+    # Get repo ID
+    repo_id = input("Enter repository ID (e.g., username/ByteDream): ").strip()
+    if not repo_id:
+        print("❌ Repository ID is required!")
+        return
+    
+    print(f"\n📤 Uploading to {repo_id}...")
     
     try:
         from bytedream.generator import ByteDreamGenerator
         
-        print("\nInitializing generator...")
-        generator = ByteDreamGenerator(device="cpu")
+        # Load model
+        print("\nLoading model...")
+        generator = ByteDreamGenerator(
+            model_path="./models/bytedream",
+            config_path="config.yaml",
+            device="cpu",
+        )
         
-        print("\nModel info:")
-        info = generator.get_model_info()
-        for key, value in info.items():
-            print(f"  {key}: {value}")
+        # Upload
+        generator.push_to_hub(
+            repo_id=repo_id,
+            token=token,
+            private=False,
+            commit_message="Upload Byte Dream model",
+        )
         
-        print("\nGenerating test image...")
-        print("Prompt: A simple test pattern, geometric shapes")
+        print("\n✅ SUCCESS!")
+        print(f"\n📦 Your model is available at:")
+        print(f"https://huggingface.co/{repo_id}")
+        print(f"\nTo use this model:")
+        print(f"  python infer.py --prompt 'your prompt' --hf_repo '{repo_id}'")
+        print("="*60)
         
-        image = generator.generate(
-            prompt="A simple test pattern, geometric shapes, abstract art",
-            width=256,
-            height=256,
-            num_inference_steps=20,
-            seed=42,
+    except Exception as e:
+        print(f"\n❌ Error: {e}")
+        import traceback
+        traceback.print_exc()
+
+
+def download_from_hf():
+    """Download model from Hugging Face"""
+    print_banner("DOWNLOAD FROM HUGGING FACE HUB")
+    
+    # Get repo ID
+    repo_id = input("Enter repository ID (e.g., username/ByteDream): ").strip()
+    if not repo_id:
+        print("❌ Repository ID is required!")
+        return
+    
+    print(f"\n📥 Downloading from {repo_id}...")
+    
+    try:
+        from bytedream.generator import ByteDreamGenerator
+        
+        # Load from HF
+        generator = ByteDreamGenerator(
+            hf_repo_id=repo_id,
+            config_path="config.yaml",
+            device="cpu",
         )
         
-        output_path = Path("test_output.png")
-        image.save(output_path)
+        print("\n✅ Model loaded successfully!")
         
-        print(f"\n✓ Test successful!")
-        print(f"  Image saved to: {output_path.absolute()}")
+        # Test generation
+        test = input("\nGenerate test image? (y/n): ").strip().lower()
+        if test == 'y':
+            print("\nGenerating test image...")
+            image = generator.generate(
+                prompt="test pattern, simple colors",
+                width=256,
+                height=256,
+                num_inference_steps=10,
+            )
+            
+            output = "test_output.png"
+            image.save(output)
+            print(f"✓ Test image saved to: {output}")
         
-        return True
+        print("\nTo generate images:")
+        print(f"  python infer.py --prompt 'your prompt' --hf_repo '{repo_id}'")
+        print(f"  HF_REPO_ID={repo_id} python app.py")
+        print("="*60)
         
     except Exception as e:
-        print(f"\n✗ Error: {e}")
-        print("\nThe model needs to be trained or pretrained weights downloaded.")
-        return False
+        print(f"\n❌ Error: {e}")
+        import traceback
+        traceback.print_exc()
 
 
-def download_pretrained():
-    """Download pretrained model from Hugging Face"""
-    print("\n" + "="*60)
-    print("Downloading Pretrained Model")
-    print("="*60)
-    
-    print("\nTo download a pretrained model:")
-    print("1. Visit https://huggingface.co/models")
-    print("2. Search for 'stable-diffusion' or similar")
-    print("3. Download using:")
-    print("\n   from huggingface_hub import snapshot_download")
-    print("   snapshot_download(repo_id='username/model', local_dir='./models/bytedream')")
-    print("\nOr train your own model with:")
-    print("   python train.py --train_data ./dataset --output_dir ./models/bytedream")
+def test_local_model():
+    """Test local model"""
+    print_banner("TEST LOCAL MODEL")
+    
+    if not check_model_exists():
+        return
+    
+    print("Loading local model...")
+    
+    try:
+        from bytedream.generator import ByteDreamGenerator
+        
+        generator = ByteDreamGenerator(
+            model_path="./models/bytedream",
+            config_path="config.yaml",
+            device="cpu",
+        )
+        
+        print("\n✅ Model loaded successfully!")
+        
+        # Generate test image
+        print("\nGenerating test image...")
+        image = generator.generate(
+            prompt="test pattern, simple colors",
+            width=256,
+            height=256,
+            num_inference_steps=10,
+        )
+        
+        output = "test_output.png"
+        image.save(output)
+        print(f"✓ Test image saved to: {output}")
+        
+        print("\nModel ready for upload!")
+        print("To upload: python quick_start.py upload")
+        print("="*60)
+        
+    except Exception as e:
+        print(f"\n❌ Error: {e}")
+        import traceback
+        traceback.print_exc()
 
 
 def main():
-    print("="*60)
-    print("Byte Dream - Quick Start")
-    print("="*60)
-    
-    # Check requirements
-    if not check_requirements():
-        print("\n⚠ Please install requirements first")
-        sys.exit(1)
-    
-    # Test model
-    if test_model():
-        print("\n✓ Byte Dream is ready to use!")
-        print("\nNext steps:")
-        print("  - Run: python infer.py --prompt 'Your prompt here'")
-        print("  - Run: python app.py (for web interface)")
-        print("  - Run: python main.py --demo (for demo)")
+    """Main function"""
+    print_banner("BYTE DREAM - QUICK START")
+    
+    print("What would you like to do?")
+    print("1. Upload model to Hugging Face")
+    print("2. Download model from Hugging Face")
+    print("3. Test local model")
+    print("4. Exit")
+    print()
+    
+    choice = input("Enter choice (1-4): ").strip()
+    
+    if choice == "1":
+        upload_to_hf()
+    elif choice == "2":
+        download_from_hf()
+    elif choice == "3":
+        test_local_model()
+    elif choice == "4":
+        print("\nGoodbye!")
+        return
     else:
-        download_pretrained()
+        print("❌ Invalid choice!")
+        return
+    
+    print("\nDone!")
 
 
 if __name__ == "__main__":
-    main()
+    try:
+        main()
+    except KeyboardInterrupt:
+        print("\n\nInterrupted!")
+        sys.exit(0)