Update app.py (#26)

- Update app.py (c22eda8dbeff393998ea7a5771fb822cda465ee3)
- Update app.py (0b8a50cce4668cb4992119e0e9fa210d81703402)

app.py

@@ -6,9 +6,9 @@ import shutil
 from huggingface_hub import HfApi, create_repo
 from typing import Union
 
-def download_file(digest, image):
+def download_file(media_type, digest, image):
     url = f"https://registry.ollama.ai/v2/library/{image}/blobs/{digest}"
-    file_name = f"blobs/{digest}"
+    file_name = f"blobs/{media_type.split(".")[-1]}"
 
     # Create the directory if it doesn't exist
     os.makedirs(os.path.dirname(file_name), exist_ok=True)
@@ -65,17 +65,18 @@ def process_image_tag(image_tag, repo_id, oauth_token: Union[gr.OAuthToken, None
         with open(manifest_file_path, 'w') as f:
             json.dump(manifest_json, f)
     
-        # Extract the digest values from the JSON
-        digests = [layer['digest'] for layer in manifest_json.get('layers', [])]
-    
-        # Download each file
-        for digest in digests:
-            download_file(digest, image)
+        # Extract the mediaType and digest values from the JSON
+        layers = manifest_json.get('layers', [])
+        for layer in layers:
+            media_type = layer['mediaType']
+            digest = layer['digest']
+            download_file(media_type, digest, image)
     
         # Download the config file
+        config_media_type = manifest_json.get('config', {}).get('mediaType')
         config_digest = manifest_json.get('config', {}).get('digest')
-        if config_digest:
-            download_file(config_digest, image)
+        if config_media_type and config_digest:
+            download_file(config_media_type, config_digest, image)
 
         # Upload to Hugging Face Hub
         upload_result, repo_path = upload_to_huggingface(repo_id, 'blobs', oauth_token)