Update app.py

app.py

@@ -2,6 +2,7 @@ import gradio as gr
 import os
 import requests
 import json
+import shutil
 from huggingface_hub import HfApi, create_repo
 
 HF_TOKEN = os.environ.get("HF_TOKEN")
@@ -30,9 +31,9 @@ def fetch_manifest(image, tag):
     else:
         return None
 
-def upload_to_huggingface(repo_id, folder_path):
-    api = HfApi(token=HF_TOKEN)
-    repo_path = api.create_repo(repo_id, "model", exist_ok=True)
+def upload_to_huggingface(repo_id, folder_path, token=token):
+    api = HfApi(token=token)
+    repo_path = api.create_repo(repo_id=repo_id, repo_type="model", exist_ok=True)
     print(f"Repo created {repo_path}")
     try:
         api.upload_folder(
@@ -44,9 +45,11 @@ def upload_to_huggingface(repo_id, folder_path):
     except Exception as e:
         return f"Upload failed: {str(e)}"
 
-def process_image_tag(image_tag, repo_id):
+def process_image_tag(image_tag, repo_id, oauth_token: Union[gr.OAuthToken, None]):
     # Extract image and tag from the input
     try:
+        token = oauth_token.token if oauth_token else None
+        
         image, tag = image_tag.split(':')
     
         # Fetch the manifest JSON
@@ -73,23 +76,28 @@ def process_image_tag(image_tag, repo_id):
             download_file(config_digest, image)
     
         # Upload to Hugging Face Hub
-        upload_result = upload_to_huggingface(repo_id, 'blobs')
+        upload_result = upload_to_huggingface(repo_id, 'blobs/*', token=token)
     
         # Delete the blobs folder
-        os.rmtree('blobs')
+        shutil.rmtree('blobs')
         return f"Successfully fetched and downloaded files for {image}:{tag}\n{upload_result}\nBlobs folder deleted"
     except Exception as e:
-        os.rmtree('blobs', ignore_errors=True)
+        shutil.rmtree('blobs', ignore_errors=True)
         return f"Error found: {str(e)}"
 
+css = """
+.main_ui_logged_out{opacity: 0.3; pointer-events: none}
+"""
+
 # Create the Gradio interface using gr.Blocks
-with gr.Blocks() as demo:
+with gr.Blocks(css=css) as demo:
     gr.Markdown("# Ollama <> HF Hub 🤝")
     gr.Markdown("Enter the image and tag to download the corresponding files from the Ollama registry and upload them to the Hugging Face Hub.")
     
-    with gr.Row():
-        image_tag_input = gr.Textbox(placeholder="Enter Ollama ID", label="Image and Tag")
-        repo_id_input = gr.Textbox(placeholder="Enter Hugging Face repo ID", label="Hugging Face Repo ID")
+    gr.LoginButton()
+    
+    image_tag_input = gr.Textbox(placeholder="Enter Ollama ID", label="Image and Tag")
+    repo_id_input = gr.Textbox(placeholder="Enter Hugging Face repo ID", label="Hugging Face Repo ID")
     
     result_output = gr.Textbox(label="Result")