Update app.py

app.py

@@ -602,23 +602,54 @@ def check_existing_model(model_path: str) -> bool:
     return all(f in model_files for f in required_files) and has_weights
 
 async def download_model_files():
-    """Downloads the model files using git clone from Hugging Face repository"""
+    """Downloads the model files using Hugging Face Hub API"""
     try:
         print(f"[INFO] Processing model from {Settings.MODEL_REPO}...")
         
+        # Install required packages if not present
+        required_packages = ["huggingface_hub", "requests", "tqdm"]
+        for package in required_packages:
+            try:
+                __import__(package)
+            except ImportError:
+                print(f"[INFO] Installing {package}...")
+                import subprocess
+                subprocess.check_call(["pip", "install", package])
+        
+        from huggingface_hub import hf_hub_download, snapshot_download, HfFolder
+        import requests
+        from tqdm import tqdm
+        
         # Create models directory
         models_dir = os.path.join(os.getcwd(), "models")
         os.makedirs(models_dir, exist_ok=True)
         print(f"[INFO] Models directory: {models_dir}")
         
         # Get the model name from the repository URL
-        model_name = Settings.MODEL_REPO.split('/')[-1]
+        repo_id = "/".join(Settings.MODEL_REPO.split('/')[-2:])  # e.g., "facebook/opt-125m"
+        model_name = repo_id.split('/')[-1]
         
         # Create versioned model directory
         version = get_next_model_version(models_dir, model_name)
         model_base_dir = os.path.join(models_dir, model_name)
         model_version_dir = os.path.join(model_base_dir, f"v{version}")
         
+        # Function to download file with progress bar
+        def download_file(url, filename):
+            response = requests.get(url, stream=True)
+            total_size = int(response.headers.get('content-length', 0))
+            
+            with open(filename, 'wb') as f, tqdm(
+                desc=os.path.basename(filename),
+                total=total_size,
+                unit='iB',
+                unit_scale=True,
+                unit_divisor=1024,
+            ) as pbar:
+                for data in response.iter_content(chunk_size=1024):
+                    size = f.write(data)
+                    pbar.update(size)
+        
         # Check if previous version exists and is valid
         if version > 1:
             prev_version_dir = os.path.join(model_base_dir, f"v{version-1}")
@@ -627,21 +658,53 @@ async def download_model_files():
                 model_path = prev_version_dir
                 state.is_model_loaded = True
             else:
-                # Clone new version if previous is invalid or incomplete
+                # Download new version
                 os.makedirs(model_version_dir, exist_ok=True)
-                success = clone_repository(Settings.MODEL_REPO, model_version_dir)
-                if not success:
-                    raise Exception("Failed to clone repository")
                 model_path = model_version_dir
-                print(f"[INFO] Successfully cloned model to {model_path}")
         else:
             # First time download
             os.makedirs(model_version_dir, exist_ok=True)
-            success = clone_repository(Settings.MODEL_REPO, model_version_dir)
-            if not success:
-                raise Exception("Failed to clone repository")
             model_path = model_version_dir
-            print(f"[INFO] Successfully cloned model to {model_path}")
+        
+        if not state.is_model_loaded:
+            try:
+                print(f"[INFO] Downloading model files from {repo_id}...")
+                
+                # First download config and other small files
+                config_files = ["config.json", "tokenizer_config.json", "vocab.json", "generation_config.json"]
+                for filename in config_files:
+                    try:
+                        file_path = hf_hub_download(
+                            repo_id=repo_id,
+                            filename=filename,
+                            local_dir=model_path,
+                            force_download=True
+                        )
+                        print(f"[INFO] Downloaded {filename}")
+                    except Exception as e:
+                        print(f"[WARN] Could not download {filename}: {str(e)}")
+                
+                # Then download the model weights
+                print("[INFO] Downloading model weights (this may take a while)...")
+                for weight_file in ["pytorch_model.bin", "model.safetensors"]:
+                    try:
+                        file_path = hf_hub_download(
+                            repo_id=repo_id,
+                            filename=weight_file,
+                            local_dir=model_path,
+                            force_download=True
+                        )
+                        print(f"[INFO] Successfully downloaded {weight_file}")
+                        break  # Stop after first successful weight file download
+                    except Exception as e:
+                        print(f"[WARN] Could not download {weight_file}: {str(e)}")
+                        continue
+                
+                print(f"[INFO] All files downloaded to {model_path}")
+                state.is_model_loaded = True
+                
+            except Exception as e:
+                raise Exception(f"Failed to download model files: {str(e)}")
         
         # Set model paths in state
         state.model_path = model_path
@@ -1091,7 +1154,7 @@ if __name__ == "__main__":
     print(f"[INFO] API Documentation available at http://localhost:{port}/docs")
     
     uvicorn.run(
-        "app:app",
+        "controller_server_new:app",
         host="0.0.0.0",
         port=port,
         reload=False