Update app.py

app.py

@@ -1,65 +1,100 @@
-import os
-import shutil
 import gradio as gr
-from huggingface_hub import HfApi, snapshot_download
+from huggingface_hub import HfApi, CommitOperationDelete
 
-def clone_repository(source_repo, target_repo, hf_token):
+def fetch_files(source_repo, hf_token):
+    # Input validation
     if not hf_token.strip():
-        return "Error: Hugging Face Access Token is required."
-    if not source_repo.strip() or not target_repo.strip():
-        return "Error: Both source and target repository names are required."
+        yield gr.update(), gr.update(), "Error: Hugging Face Access Token is required.", []
+        return
+    if not source_repo.strip():
+        yield gr.update(), gr.update(), "Error: Source repository name is required.", []
+        return
     
     api = HfApi(token=hf_token.strip())
-    local_dir = "temp_clone_dir"
     
     try:
-        # 1. Create the target repository on your account
-        yield f"Status: Creating target repository '{target_repo}'..."
-        api.create_repo(repo_id=target_repo.strip(), repo_type="model", exist_ok=True)
+        # Initial status update before fetching
+        yield gr.update(), gr.update(), f"Status: Fetching file list from '{source_repo}'...", []
         
-        # 2. Clean up any leftover temporary directories from previous runs
-        if os.path.exists(local_dir):
-            shutil.rmtree(local_dir)
+        # Fetch the flat list of files from the repository
+        files = api.list_repo_files(repo_id=source_repo.strip(), repo_type="model")
+        files.sort()
+        
+        # Build a visual Linux-like tree representation for the CheckboxGroup
+        choices = []
+        for f in files:
+            parts = f.split('/')
+            if len(parts) == 1:
+                display = f"📄 {f}"
+            else:
+                folder_path = "/".join(parts[:-1])
+                display = f"📁 {folder_path}/ ├── 📄 {parts[-1]}"
+            
+            # Gradio 4.0+ allows (display_name, value) tuples for CheckboxGroup choices
+            choices.append((display, f))
             
-        # 3. Download the source repository files
-        yield f"Status: Downloading files from '{source_repo}' (this may take some time depending on the model size)..."
-        downloaded_dir = snapshot_download(
-            repo_id=source_repo.strip(),
-            local_dir=local_dir,
-            token=hf_token.strip(),
-            max_workers=4
+        yield (
+            gr.update(visible=True), 
+            gr.update(choices=choices, value=files), 
+            f"Success: Fetched {len(files)} files. Uncheck the ones you do not want to copy.",
+            files
         )
         
-        # 4. Upload the files to your target repository
-        yield f"Status: Uploading files to '{target_repo}'..."
-        api.upload_folder(
-            folder_path=downloaded_dir,
-            repo_id=target_repo.strip(),
+    except Exception as e:
+        yield gr.update(visible=False), gr.update(), f"An error occurred: {str(e)}", []
+
+
+def clone_repository(source_repo, target_repo, hf_token, selected_files, all_files):
+    if not target_repo.strip():
+        yield "Error: Target repository name is required."
+        return
+    
+    api = HfApi(token=hf_token.strip())
+    
+    try:
+        yield f"Status: Initiating server-side clone of '{source_repo}' to '{target_repo}'..."
+        
+        # 1. Duplicate the entire repo first (Server-side, almost instant, 0 disk space)
+        api.duplicate_repo(
+            from_id=source_repo.strip(),
+            to_id=target_repo.strip(),
             repo_type="model"
         )
         
-        # 5. Clean up local storage
-        shutil.rmtree(local_dir)
+        # 2. Determine which files to delete based on what the user unchecked
+        files_to_delete = set(all_files) - set(selected_files)
         
-        yield f"Success! Model successfully cloned to: https://huggingface.co/{target_repo.strip()}"
+        if files_to_delete:
+            yield f"Status: Deleting {len(files_to_delete)} unselected files from '{target_repo}'..."
+            
+            # Create a list of deletion operations
+            operations = [CommitOperationDelete(path=f) for f in files_to_delete]
+            
+            # Execute deletion in a single commit natively on HF servers
+            api.create_commit(
+                repo_id=target_repo.strip(),
+                repo_type="model",
+                operations=operations,
+                commit_message="Clean up: Remove unselected files during selective clone"
+            )
+        
+        yield f"Success! Selective clone completed. View your model here: https://huggingface.co/{target_repo.strip()}"
         
     except Exception as e:
-        # Ensure cleanup happens even if an error occurs
-        if os.path.exists(local_dir):
-            shutil.rmtree(local_dir)
         yield f"An error occurred: {str(e)}"
 
+
 # Build the Gradio Interface using the .Soft theme
-with gr.Blocks(theme=gr.themes.Soft(), title="HF Model Cloner") as demo:
-    gr.Markdown("# 📑 Hugging Face Model Cloner")
-    gr.Markdown("Clone any public or accessible model repository to your own Hugging Face account under a custom name.")
+with gr.Blocks(theme=gr.themes.Soft(), title="Selective HF Model Cloner") as demo:
+    gr.Markdown("# 📑 Selective Hugging Face Model Cloner")
+    gr.Markdown("Fetch files from a model repository, uncheck what you don't need, and clone via high-speed server-side duplication.")
     
     with gr.Row():
         hf_token = gr.Textbox(
             label="Hugging Face Access Token", 
             placeholder="hf_...", 
             type="password",
-            info="Requires 'Write' permissions to create and upload to your repository."
+            info="Requires 'Write' permissions."
         )
     
     with gr.Row():
@@ -67,18 +102,37 @@ with gr.Blocks(theme=gr.themes.Soft(), title="HF Model Cloner") as demo:
             label="Source Model Repository", 
             placeholder="e.g., meta-llama/Llama-3.2-1B"
         )
+        fetch_btn = gr.Button("1. Fetch Files", variant="secondary")
+        
+    # We use a Gradio State to quietly keep track of the complete list of files fetched
+    all_files_state = gr.State([])
+    
+    # This group remains hidden until files are successfully fetched
+    with gr.Group(visible=False) as selection_group:
+        gr.Markdown("### 🗂️ Repository Files")
+        gr.Markdown("*All files are selected by default. Uncheck any files you do **not** want to carry over.*")
+        
+        file_checkboxes = gr.CheckboxGroup(label="Repository Tree", interactive=True)
+        
         target_repo = gr.Textbox(
             label="Target Model Repository", 
             placeholder="e.g., your-username/my-custom-llama"
         )
+        clone_btn = gr.Button("2. Clone Selected Files", variant="primary")
         
-    clone_btn = gr.Button("Start Cloning Process", variant="primary")
     output_logs = gr.Textbox(label="Execution Status / Logs", interactive=False)
     
-    # Use streaming output (yield) to show live progress updates
+    # Trigger Step 1: File Fetching
+    fetch_btn.click(
+        fn=fetch_files,
+        inputs=[source_repo, hf_token],
+        outputs=[selection_group, file_checkboxes, output_logs, all_files_state]
+    )
+    
+    # Trigger Step 2: Cloning Process
     clone_btn.click(
         fn=clone_repository,
-        inputs=[source_repo, target_repo, hf_token],
+        inputs=[source_repo, target_repo, hf_token, file_checkboxes, all_files_state],
         outputs=output_logs
     )