Spaces:
Running
Running
File size: 3,220 Bytes
38723d6 76a4a39 38723d6 c4a67c0 76a4a39 38723d6 76a4a39 38723d6 76a4a39 38723d6 76a4a39 38723d6 76a4a39 c4a67c0 76a4a39 38723d6 76a4a39 38723d6 76a4a39 38723d6 76a4a39 38723d6 76a4a39 38723d6 76a4a39 38723d6 76a4a39 38723d6 76a4a39 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 |
import gradio as gr
import os
import requests
import json
from huggingface_hub import HfApi, create_repo
HF_TOKEN = os.environ.get("HF_TOKEN")
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(token=HF_TOKEN)
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() |