Upload app.py

app.py

@@ -22,62 +22,35 @@ from pipeline_flux_controlnet_inpaint import FluxControlNetInpaintingPipeline
 import huggingface_hub
 huggingface_hub.login(os.getenv('HF_TOKEN_FLUX'))
 
+import gdown
 
-def download_release_asset(repo, release_tag, asset_name, destination, token):
+def download_from_gdrive(file_id, destination):
     """
-    Download a file from a GitHub Release.
+    Download a file from Google Drive using gdown.
+    
+    Args:
+        file_id (str): The Google Drive file ID.
+        destination (str): Local path to save the downloaded file.
     """
-    url = f"https://api.github.com/repos/{repo}/releases/tags/{release_tag}"
-    headers = {"Authorization": f"token {token}"}
-    print(f"Fetching release info from: {url}")  # Debug URL
-
-    response = requests.get(url, headers=headers)
-    print(f"Status Code: {response.status_code}")  # Debug status code
-    print(f"Response Content: {response.json()}")  # Debug API response
-
-    if response.status_code == 200:
-        release_data = response.json()
-        # Find the asset by name
-        asset = next((a for a in release_data['assets'] if a['name'] == asset_name), None)
-        if not asset:
-            raise Exception(f"Asset {asset_name} not found in release {release_tag}.")
-
-        # Download the asset
-        download_url = asset['browser_download_url']
-        print(f"Downloading asset from: {download_url}")  # Debug download URL
-
-        download_response = requests.get(download_url, headers=headers, stream=True)
-        if download_response.status_code == 200:
-            os.makedirs(os.path.dirname(destination), exist_ok=True)
-            with open(destination, 'wb') as f:
-                for chunk in download_response.iter_content(chunk_size=1024):
-                    f.write(chunk)
-            print(f"Downloaded {asset_name} to {destination}")
-        else:
-            raise Exception(f"Failed to download {asset_name}. Status code: {download_response.status_code}")
-    else:
-        raise Exception(f"Failed to fetch release {release_tag}. Status code: {response.status_code}")
-
-
-
-# Example: Define GitHub repo and file paths
-repo_name = "arj7192/NativDemoLocal"
-release_tag = "v1"
-github_files = {
-    "speech_bubble_model.pt": "speech_bubble_model.pt",
-    "craft_mlt_25k.pth": "craft_mlt_25k.pth",
-    "english_g2.pth": "english_g2.pth",
-    "korean_g2.pth": "korean_g2.pth",
-    "latin_g2.pth": "latin_g2.pth",
-    "zh_sim_g2.pth": "zh_sim_g2.pth",
+    url = f"https://drive.google.com/uc?id={file_id}"
+    gdown.download(url, destination, quiet=False)
+    print(f"File downloaded to {destination}")
+
+
+files = {
+    "speech_bubble_model": "speech_bubble_model.pt",
+    "craft_mlt_25k": "craft_mlt_25k.pth",
+    "english_g2": "english_g2.pth",
+    "korean_g2": "korean_g2.pth",
+    "latin_g2": "latin_g2.pth",
+    "zh_sim_g2": "zh_sim_g2.pth",
 }
 
 token = os.getenv("HF_GITHUB_TOKEN")
 
 # Download each file
-for filename, repo_path in github_files.items():
-    destination_path = f"{filename}"
-    download_release_asset(repo_name, release_tag, filename, destination_path, token)
+for filename, destination_path in files.items():
+    download_from_gdrive(os.getenv(filename), destination_path)
 
 bubble_detection_model = YOLO("speech_bubble_model.pt")