Upload: app.py

app.py

@@ -0,0 +1,622 @@
+import gradio as gr
+import os
+import logging
+import asyncio
+import tempfile
+import git
+import shutil
+import zipfile
+from pathlib import Path
+from typing import Tuple, Optional
+from huggingface_hub import login, HfApi, create_repo, upload_file
+from urllib.parse import urlparse
+from functools import partial
+from datetime import datetime
+
+# Setup logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(levelname)s - [%(filename)s:%(lineno)d] - %(message)s'
+)
+logger = logging.getLogger(__name__)
+
+class HuggingFaceManager:
+    def __init__(self):
+        self.api = HfApi()
+        self.token: Optional[str] = None
+        self.temp_dir = Path(tempfile.mkdtemp())
+        
+    def cleanup(self):
+        """Clean up temporary directory"""
+        if self.temp_dir.exists():
+            shutil.rmtree(self.temp_dir, ignore_errors=True)
+            
+    def validate_repo_name(self, repo_name: str) -> bool:
+        """Validate repository name format"""
+        if not repo_name or '/' not in repo_name:
+            raise ValueError("Repository name must be in format 'username/repo-name'")
+        username, repo = repo_name.split('/', 1)
+        return all(name.strip() for name in [username, repo])
+
+    def validate_url(self, url: str) -> bool:
+        """Validate URL format"""
+        try:
+            result = urlparse(url)
+            return all([result.scheme, result.netloc])
+        except Exception:
+            return False
+
+    async def create_zip_archive(
+        self,
+        source_path: Path,
+        password: Optional[str] = None,
+        progress_callback=None
+    ) -> Optional[Path]:
+        """Create a zip archive of the folder, optionally with password protection"""
+        try:
+            if not source_path.exists():
+                raise ValueError(f"Source path {source_path} does not exist")
+            timestamp = datetime.now().strftime("%S_%M_%H_%d_%m_%Y")
+            zip_path = self.temp_dir / f"archive_{timestamp}.zip"  # Tạo tên file zip
+            
+            # Count total files for progress
+            total_files = sum(1 for _ in source_path.rglob('*') if _.is_file())
+            processed_files = 0
+
+            # Create zip archive
+            if password:
+                # Using external zip command for password protection
+                files_list = [
+                    str(f.relative_to(source_path))
+                    for f in source_path.rglob('*')
+                    if f.is_file()
+                ]
+                
+                # Create temporary file listing
+                list_file = self.temp_dir / "files.txt"
+                list_file.write_text("\n".join(files_list))
+                
+                # Use zip command with password
+                process = await asyncio.create_subprocess_exec(
+                    'zip', '-r', '-P', password, str(zip_path), '.',
+                    cwd=str(source_path),
+                    stdout=asyncio.subprocess.PIPE,
+                    stderr=asyncio.subprocess.PIPE
+                )
+                
+                await process.communicate()
+                if process.returncode != 0:
+                    raise Exception("Failed to create password-protected zip")
+                
+                if progress_callback:
+                    progress_callback(1.0)  # Complete
+            else:
+                # Use Python's zipfile for non-password zip
+                with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
+                    for file_path in source_path.rglob('*'):
+                        if file_path.is_file():
+                            relative_path = file_path.relative_to(source_path)
+                            zipf.write(file_path, relative_path)
+                            processed_files += 1
+                            if progress_callback:
+                                progress_callback(processed_files / total_files)
+
+            return zip_path
+        except Exception as e:
+            logger.error(f"Failed to create zip archive: {e}")
+            return None
+            
+    async def login_and_validate(self, token: str) -> str:
+        """Login to Hugging Face and validate token"""
+        if not token.strip():
+            return "Error: Token cannot be empty"
+        
+        try:
+            loop = asyncio.get_event_loop()
+            await loop.run_in_executor(None, login, token)
+            await loop.run_in_executor(None, self.api.whoami)
+            
+            self.token = token
+            return "Login successful!"
+        except Exception as e:
+            logger.error(f"Login failed: {e}")
+            return f"Login failed: {str(e)}"
+
+    async def download_from_url(
+        self, 
+        url: str, 
+        download_type: str,
+        progress: Optional[gr.Progress] = None
+    ) -> Tuple[str, Optional[str]]:
+        """Download content from URL using wget or git"""
+        if not self.validate_url(url):
+            return "Invalid URL format!", None
+
+        output_path = self.temp_dir / "downloaded_content"
+        if output_path.exists():
+            shutil.rmtree(output_path)
+        output_path.mkdir(parents=True)
+
+        try:
+            if download_type == "wget":
+                if progress:
+                    progress(0, desc="Downloading with wget...")
+                
+                process = await asyncio.create_subprocess_exec(
+                    'wget', '-P', str(output_path), url,
+                    stdout=asyncio.subprocess.PIPE,
+                    stderr=asyncio.subprocess.PIPE
+                )
+                
+                try:
+                    stdout, stderr = await asyncio.wait_for(process.communicate(), timeout=300)
+                    if process.returncode != 0:
+                        raise Exception(f"wget failed: {stderr.decode()}")
+                except asyncio.TimeoutError:
+                    raise Exception("Download timed out after 5 minutes")
+
+            elif download_type == "git":
+                if progress:
+                    progress(0, desc="Cloning git repository...")
+                
+                await asyncio.get_event_loop().run_in_executor(
+                    None,
+                    partial(git.Repo.clone_from, url, str(output_path), depth=1)
+                )
+
+            if progress:
+                progress(1, desc="Download complete")
+            
+            return "Download successful!", str(output_path)
+
+        except Exception as e:
+            logger.error(f"Download failed: {e}")
+            if output_path.exists():
+                shutil.rmtree(output_path)
+            return f"Download failed: {str(e)}", None
+
+    async def create_repository(
+        self, 
+        repo_name: str, 
+        repo_type: str, 
+        is_private: bool
+    ) -> str:
+        """Create a new repository on Hugging Face"""
+        if not self.token:
+            return "Please login first!"
+
+        try:
+            self.validate_repo_name(repo_name)
+            
+            await asyncio.get_event_loop().run_in_executor(
+                None,
+                partial(
+                    create_repo,
+                    repo_name,
+                    private=is_private,
+                    repo_type=repo_type,
+                    exist_ok=True
+                )
+            )
+            
+            return f"Repository '{repo_name}' created successfully!"
+        except Exception as e:
+            logger.error(f"Failed to create repository: {e}")
+            return f"Failed to create repository: {str(e)}"
+
+    async def upload_file_worker(
+        self, 
+        file_path: Path, 
+        relative_path: Path, 
+        repo_name: str, 
+        repo_type: str,
+        progress_callback=None
+    ) -> Optional[str]:
+        """Upload a single file to repository"""
+        try:
+            await asyncio.get_event_loop().run_in_executor(
+                None,
+                partial(
+                    upload_file,
+                    path_or_fileobj=str(file_path),
+                    path_in_repo=str(relative_path),
+                    repo_id=repo_name,
+                    repo_type=repo_type,
+                    commit_message=f"Upload: {relative_path}"
+                )
+            )
+            if progress_callback:
+                progress_callback()
+            return str(relative_path)
+        except Exception as e:
+            logger.error(f"Error uploading {relative_path}: {e}")
+            return None
+
+    async def upload_folder(
+        self, 
+        folder_path: str, 
+        repo_name: str, 
+        repo_type: str,
+        should_zip: bool = False,
+        zip_password: Optional[str] = None,
+        progress: Optional[gr.Progress] = None
+    ) -> str:
+        """Upload entire folder to repository, optionally as a zip archive"""
+        if not self.token:
+            return "Please login first!"
+
+        folder_path = Path(folder_path)
+        if not folder_path.exists():
+            return f"Folder {folder_path} does not exist!"
+
+        try:
+            self.validate_repo_name(repo_name)
+            
+            if should_zip:
+                if progress:
+                    progress(0, desc="Creating zip archive...")
+                
+                zip_path = await self.create_zip_archive(
+                    folder_path,
+                    zip_password,
+                    lambda p: progress(p * 0.5) if progress else None
+                )
+                
+                if not zip_path:
+                    return "Failed to create zip archive"
+
+                if progress:
+                    progress(0.5, desc="Uploading zip file...")
+
+                # Upload single zip file
+                await self.upload_file_worker(
+                    zip_path,
+                    Path(zip_path.name),
+                    repo_name,
+                    repo_type,
+                    lambda: progress(1.0) if progress else None
+                )
+                
+                return "Zip file uploaded successfully!"
+            else:
+                # Original folder upload logic
+                files_to_upload = []
+                for root, dirs, files in os.walk(folder_path):
+                    if '.git' in dirs:
+                        dirs.remove('.git')
+                    
+                    for file in files:
+                        file_path = Path(root) / file
+                        relative_path = file_path.relative_to(folder_path)
+                        files_to_upload.append((file_path, relative_path))
+
+                if not files_to_upload:
+                    return "No files found to upload"
+
+                if progress:
+                    progress(0, desc=f"Uploading {len(files_to_upload)} files...")
+
+                # Track progress
+                uploaded_count = 0
+                def update_progress():
+                    nonlocal uploaded_count
+                    uploaded_count += 1
+                    if progress:
+                        progress(uploaded_count / len(files_to_upload))
+
+                # Upload files in parallel
+                tasks = [
+                    self.upload_file_worker(
+                        file_path, 
+                        relative_path, 
+                        repo_name, 
+                        repo_type, 
+                        update_progress
+                    )
+                    for file_path, relative_path in files_to_upload
+                ]
+                
+                results = await asyncio.gather(*tasks, return_exceptions=True)
+                
+                successful = sum(1 for r in results if r is not None)
+                failed = len(files_to_upload) - successful
+                
+                return f"Uploaded {successful} files successfully. Failed: {failed}"
+
+        except Exception as e:
+            logger.error(f"Upload failed: {e}")
+            return f"Upload failed: {str(e)}"
+        finally:
+            # Cleanup temporary zip file if created
+            if should_zip:
+                zip_path = self.temp_dir / "archive.zip"
+                if zip_path.exists():
+                    zip_path.unlink()
+
+    async def delete_files(
+        self, 
+        folder_path: str, 
+        repo_name: str, 
+        repo_type: str, 
+        dry_run: bool = True,
+        progress: Optional[gr.Progress] = None
+    ) -> str:
+        """Delete files from repository based on local folder structure"""
+        if not self.token:
+            return "Please login first!"
+
+        try:
+            self.validate_repo_name(repo_name)
+            
+            folder_path = Path(folder_path)
+            if not folder_path.exists():
+                return f"Folder {folder_path} does not exist!"
+
+            # Get repository files
+            existing_files = await asyncio.get_event_loop().run_in_executor(
+                None,
+                partial(self.api.list_repo_files, repo_id=repo_name, repo_type=repo_type)
+            )
+
+            # Find files to delete
+            files_to_delete = []
+            for root, dirs, files in os.walk(folder_path):
+                if '.git' in dirs:
+                    dirs.remove('.git')
+                
+                for file in files:
+                    file_path = Path(root) / file
+                    relative_path = str(file_path.relative_to(folder_path))
+                    if relative_path in existing_files:
+                        files_to_delete.append(relative_path)
+
+            if not files_to_delete:
+                return "No matching files found to delete"
+
+            if dry_run:
+                return f"Dry run: Would delete {len(files_to_delete)} files"
+
+            # Perform deletion
+            if progress:
+                progress(0, desc=f"Deleting {len(files_to_delete)} files...")
+            
+            deleted_count = 0
+            for file_path in files_to_delete:
+                try:
+                    await asyncio.get_event_loop().run_in_executor(
+                        None,
+                        partial(
+                            self.api.delete_file,
+                            repo_id=repo_name,
+                            path_in_repo=file_path,
+                            repo_type=repo_type
+                        )
+                    )
+                    deleted_count += 1
+                    if progress:
+                        progress(deleted_count / len(files_to_delete))
+                except Exception as e:
+                    logger.error(f"Error deleting {file_path}: {e}")
+
+            return f"Successfully deleted {deleted_count} out of {len(files_to_delete)} files"
+
+        except Exception as e:
+            logger.error(f"Delete operation failed: {e}")
+            return f"Delete operation failed: {str(e)}"
+
+def create_interface() -> gr.Blocks:
+    """Create Gradio interface"""
+    manager = HuggingFaceManager()
+    
+    with gr.Blocks(title="Hugging Face Repository Manager") as app:
+        gr.Markdown("# Hugging Face Repository Manager")
+        
+        # Login Tab
+        with gr.Tab("Login"):
+            token_input = gr.Textbox(
+                label="Hugging Face Token",
+                type="password",
+                placeholder="Enter your Hugging Face token"
+            )
+            login_btn = gr.Button("Login", variant="primary")
+            login_output = gr.Textbox(label="Login Status", interactive=False)
+            
+            login_btn.click(
+                fn=lambda x: asyncio.run(manager.login_and_validate(x)),
+                inputs=[token_input],
+                outputs=[login_output]
+            )
+        
+        # Download Tab
+        with gr.Tab("Download"):
+            download_url = gr.Textbox(
+                label="URL",
+                placeholder="Enter URL to download"
+            )
+            download_type = gr.Radio(
+                choices=["wget", "git"],
+                label="Download Type",
+                value="wget"
+            )
+            download_btn = gr.Button("Download", variant="primary")
+            download_output = gr.Textbox(label="Status", interactive=False)
+            download_path = gr.Textbox(
+                label="Download Path",
+                interactive=False,
+                visible=False
+            )
+            
+            download_btn.click(
+                fn=lambda x, y: asyncio.run(manager.download_from_url(x, y)),
+                inputs=[download_url, download_type],
+                outputs=[download_output, download_path]
+            )
+        
+        # Create Repository Tab
+        with gr.Tab("Create Repository"):
+            repo_name = gr.Textbox(
+                label="Repository Name",
+                placeholder="username/repo-name"
+            )
+            repo_type = gr.Radio(
+                choices=["model", "dataset", "space"],
+                label="Repository Type",
+                value="model"
+            )
+            is_private = gr.Checkbox(label="Private Repository", value=True)
+            create_btn = gr.Button("Create Repository", variant="primary")
+            create_output = gr.Textbox(label="Status", interactive=False)
+            
+            create_btn.click(
+                fn=lambda x, y, z: asyncio.run(manager.create_repository(x, y, z)),
+                inputs=[repo_name, repo_type, is_private],
+                outputs=[create_output]
+            )
+        
+        # Upload Tab
+        with gr.Tab("Upload"):
+            with gr.Row():
+                upload_folder_path = gr.Textbox(
+                    label="Local Folder Path",
+                    placeholder="Path to local folder"
+                )
+                use_downloaded = gr.Checkbox(
+                    label="Use Downloaded Content",
+                    value=False
+                )
+            
+            upload_repo_name = gr.Textbox(
+                label="Repository Name",
+                placeholder="username/repo-name"
+            )
+            upload_repo_type = gr.Radio(
+                choices=["model", "dataset", "space"],
+                label="Repository Type",
+                value="model"
+            )
+            
+            with gr.Row():
+                should_zip = gr.Checkbox(
+                    label="Create Zip Archive",
+                    value=False
+                )
+                zip_password = gr.Textbox(
+                    label="Zip Password (Optional)",
+                    type="password",
+                    placeholder="Leave empty for no password",
+                    visible=False
+                )
+            
+            # Show/hide password field based on zip checkbox
+            should_zip.change(
+                fn=lambda x: gr.update(visible=x),
+                inputs=[should_zip],
+                outputs=[zip_password]
+            )
+            
+            upload_btn = gr.Button("Upload Files", variant="primary")
+            upload_output = gr.Textbox(label="Status", interactive=False)
+            
+            def prepare_upload_path(folder_path, use_downloaded, downloaded_path):
+                return downloaded_path if use_downloaded and downloaded_path else folder_path
+            
+            upload_btn.click(
+                fn=lambda *args: asyncio.run(
+                    manager.upload_folder(
+                        prepare_upload_path(args[0], args[1], args[2]),
+                        args[3],
+                        args[4],
+                        args[5],
+                        args[6] if args[6] and args[6].strip() else None
+                    )
+                ),
+                inputs=[
+                    upload_folder_path,
+                    use_downloaded,
+                    download_path,
+                    upload_repo_name,
+                    upload_repo_type,
+                    should_zip,
+                    zip_password
+                ],
+                outputs=[upload_output]
+            )
+        
+        # Delete Tab
+        with gr.Tab("Delete"):
+            delete_folder_path = gr.Textbox(
+                label="Local Folder Path",
+                placeholder="Path to local folder for reference"
+            )
+            delete_repo_name = gr.Textbox(
+                label="Repository Name",
+                placeholder="username/repo-name"
+            )
+            delete_repo_type = gr.Radio(
+                choices=["model", "dataset", "space"],
+                label="Repository Type",
+                value="model"
+            )
+            dry_run = gr.Checkbox(
+                label="Dry Run",
+                value=True,
+                info="Preview changes without making them"
+            )
+            delete_btn = gr.Button("Delete Files", variant="secondary")
+            delete_output = gr.Textbox(label="Status", interactive=False)
+            
+            delete_btn.click(
+                fn=lambda *args: asyncio.run(manager.delete_files(*args)),
+                inputs=[
+                    delete_folder_path,
+                    delete_repo_name,
+                    delete_repo_type,
+                    dry_run
+                ],
+                outputs=[delete_output]
+            )
+
+        # Error handling
+        gr.Error()
+        
+        # Footer information
+        gr.Markdown("""
+        ### Information
+        - Get your token from [Hugging Face Settings](https://huggingface.co/settings/tokens)
+        - Repository names must be in format: username/repository-name
+        - For files larger than 5GB, use Git LFS
+        - Repository types:
+          - Model: For ML models and weights
+          - Dataset: For datasets and data files
+          - Space: For demos and applications
+        """)
+
+    return app
+
+def main():
+    """Application entry point"""
+    try:
+        logger.info("Starting Hugging Face Repository Manager")
+        app = create_interface()
+        
+        # Configure and launch the application
+        app.queue()
+        app.launch(
+            server_name="0.0.0.0",
+            server_port=7860,
+            share=True,
+            max_threads=2,
+            # Security configurations
+            auth=None,
+            ssl_keyfile=None,
+            ssl_certfile=None,
+            ssl_verify=True
+        )
+    except Exception as e:
+        logger.error(f"Application failed to start: {e}")
+        raise
+    finally:
+        logger.info("Shutting down Hugging Face Repository Manager")
+
+if __name__ == "__main__":
+    main()