Spaces:
Running
Running
import gradio as gr | |
import os | |
import requests | |
import json | |
from huggingface_hub import HfApi, create_repo | |
def download_file(digest, image): | |
url = f"https://registry.ollama.ai/v2/library/{image}/blobs/{digest}" | |
file_name = f"blobs/{digest}" | |
# Create the directory if it doesn't exist | |
os.makedirs(os.path.dirname(file_name), exist_ok=True) | |
# Download the file | |
print(f"Downloading {url} to {file_name}") | |
response = requests.get(url, allow_redirects=True) | |
if response.status_code == 200: | |
with open(file_name, 'wb') as f: | |
f.write(response.content) | |
else: | |
print(f"Failed to download {url}") | |
def fetch_manifest(image, tag): | |
manifest_url = f"https://registry.ollama.ai/v2/library/{image}/manifests/{tag}" | |
response = requests.get(manifest_url) | |
if response.status_code == 200: | |
return response.json() | |
else: | |
return None | |
def upload_to_huggingface(repo_id, folder_path): | |
api = HfApi() | |
repo_path = create_repo(repo_id, "model", exists_ok=True) | |
print(f"Repo created {repo_path}") | |
try: | |
api.upload_folder( | |
folder_path=folder_path, | |
repo_id=repo_id, | |
repo_type="model", | |
) | |
return "Upload successful" | |
except Exception as e: | |
return f"Upload failed: {str(e)}" | |
def process_image_tag(image_tag, repo_id): | |
# Extract image and tag from the input | |
image, tag = image_tag.split(':') | |
# Fetch the manifest JSON | |
manifest_json = fetch_manifest(image, tag) | |
if not manifest_json or 'errors' in manifest_json: | |
return f"Failed to fetch the manifest for {image}:{tag}" | |
# Save the manifest JSON to the blobs folder | |
manifest_file_path = "blobs/manifest.json" | |
os.makedirs(os.path.dirname(manifest_file_path), exist_ok=True) | |
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) | |
# Download the config file | |
config_digest = manifest_json.get('config', {}).get('digest') | |
if config_digest: | |
download_file(config_digest, image) | |
# Upload to Hugging Face Hub | |
upload_result = upload_to_huggingface(repo_id, 'blobs/*') | |
# Delete the blobs folder | |
try: | |
os.rmtree('blobs') | |
return f"Successfully fetched and downloaded files for {image}:{tag}\n{upload_result}\nBlobs folder deleted" | |
except Exception as e: | |
return f"Failed to delete blobs folder: {str(e)}" | |
# Create the Gradio interface | |
iface = gr.Interface( | |
fn=process_image_tag, | |
inputs=[ | |
gr.Textbox(placeholder="Enter image:tag", label="Image and Tag"), | |
gr.Textbox(placeholder="Enter Hugging Face repo ID", label="Hugging Face Repo ID") | |
], | |
outputs=gr.Textbox(label="Result"), | |
title="Registry File Downloader and Uploader", | |
description="Enter the image and tag to download the corresponding files from the registry and upload them to the Hugging Face Hub." | |
) | |
# Launch the Gradio app | |
iface.launch() |