Upload bytedream/generator.py with huggingface_hub

bytedream/generator.py

@@ -24,6 +24,7 @@ class ByteDreamGenerator:
         config_path: str = "config.yaml",
         device: str = "cpu",
         use_safetensors: bool = True,
+        hf_repo_id: Optional[str] = None,
     ):
         """
         Initialize Byte Dream generator
@@ -33,17 +34,19 @@ class ByteDreamGenerator:
             config_path: Path to configuration file
             device: Device to run on (default: cpu)
             use_safetensors: Use safetensors format if available
+            hf_repo_id: Hugging Face repository ID (e.g., "username/repo")
         """
         self.device = device
         self.config_path = config_path
         self.use_safetensors = use_safetensors
+        self.hf_repo_id = hf_repo_id
         
         # Load configuration
         self.config = self._load_config(config_path)
         
         # Initialize components
         print("Initializing Byte Dream Generator...")
-        self.pipeline = self._initialize_pipeline(model_path)
+        self.pipeline = self._initialize_pipeline(model_path, hf_repo_id)
         
         # Optimize for CPU
         self._optimize_for_cpu()
@@ -92,12 +95,27 @@ class ByteDreamGenerator:
             }
         }
     
-    def _initialize_pipeline(self, model_path: Optional[str]):
+    def _initialize_pipeline(self, model_path: Optional[str], hf_repo_id: Optional[str] = None):
         """Initialize the generation pipeline"""
         from bytedream.model import create_unet, create_vae, create_text_encoder
         from bytedream.scheduler import create_scheduler
         from bytedream.pipeline import ByteDreamPipeline
         
+        # If HF repo ID is provided, try to load from Hugging Face
+        if hf_repo_id is not None:
+            print(f"Loading model from Hugging Face: {hf_repo_id}...")
+            try:
+                from bytedream.pipeline import ByteDreamPipeline
+                pipeline = ByteDreamPipeline.from_pretrained(
+                    hf_repo_id,
+                    device=self.device,
+                    dtype=torch.float32,
+                )
+                return pipeline
+            except Exception as e:
+                print(f"Error loading from Hugging Face: {e}")
+                print("Falling back to local model...")
+        
         # Create model components
         print("Creating UNet...")
         unet = create_unet(self.config)
@@ -315,3 +333,80 @@ class ByteDreamGenerator:
         if torch.cuda.is_available():
             torch.cuda.empty_cache()
         print("Memory cleared")
+    
+    def save_pretrained(self, save_directory: str):
+        """
+        Save model to directory for Hugging Face upload
+        
+        Args:
+            save_directory: Directory path to save models
+        """
+        if self.pipeline is None:
+            raise ValueError("No pipeline initialized. Cannot save.")
+        
+        return self.pipeline.save_pretrained(save_directory)
+    
+    def push_to_hub(
+        self,
+        repo_id: str,
+        token: Optional[str] = None,
+        private: bool = False,
+        commit_message: str = "Upload Byte Dream model",
+    ):
+        """
+        Push model to Hugging Face Hub
+        
+        Args:
+            repo_id: Repository ID (username/model-name)
+            token: Hugging Face API token
+            private: Whether to make repository private
+            commit_message: Commit message for the upload
+        """
+        from huggingface_hub import create_repo, HfApi
+        import tempfile
+        import shutil
+        
+        print(f"Pushing model to Hugging Face Hub: {repo_id}")
+        
+        # Create repository
+        try:
+            create_repo(
+                repo_id=repo_id,
+                token=token,
+                private=private,
+                exist_ok=True,
+                repo_type="model",
+            )
+            print("✓ Repository created/verified")
+        except Exception as e:
+            print(f"Error creating repository: {e}")
+            raise
+        
+        # Save to temporary directory
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            print(f"Saving model to temporary directory: {tmp_dir}")
+            self.save_pretrained(tmp_dir)
+            
+            # Copy config file
+            config_src = Path(self.config_path)
+            if config_src.exists():
+                config_dst = Path(tmp_dir) / "config.yaml"
+                shutil.copy2(config_src, config_dst)
+                print("✓ Config copied")
+            
+            # Upload to Hub
+            api = HfApi()
+            try:
+                api.upload_folder(
+                    folder_path=tmp_dir,
+                    repo_id=repo_id,
+                    token=token,
+                    repo_type="model",
+                    commit_message=commit_message,
+                )
+                print("✓ Model uploaded successfully!")
+                print(f"\n📦 View your model at:")
+                print(f"https://huggingface.co/{repo_id}")
+            except Exception as e:
+                print(f"Error uploading to Hub: {e}")
+                raise