File size: 3,204 Bytes
38723d6
76a4a39
 
 
 
38723d6
c60767e
 
76a4a39
 
 
38723d6
76a4a39
 
38723d6
76a4a39
 
 
 
 
 
 
 
38723d6
76a4a39
 
 
 
 
 
 
38723d6
76a4a39
c60767e
d0f893f
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")
    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()