Create app.py

app.py

@@ -0,0 +1,276 @@
+import os
+import shutil
+import subprocess
+import requests
+import re
+from urllib.parse import urlparse
+from typing import Optional, List
+from pathlib import Path
+
+from fastapi import FastAPI, HTTPException, BackgroundTasks
+from pydantic import BaseModel
+from huggingface_hub import HfApi
+from pySmartDL import SmartDL
+
+# --- dependency check ---
+try:
+    subprocess.run(["pip", "install", "--upgrade", "gdown"], check=True, capture_output=True, text=True)
+except subprocess.CalledProcessError as e:
+    print(f"Error occurred installing gdown: {e.stderr}")
+
+# --- Configuration ---
+DIR = "download"
+os.makedirs(DIR, exist_ok=True)
+
+# Helper to get token (checks env var first)
+def get_token():
+    return os.getenv('HF_token')
+
+def get_civit_key():
+    return os.getenv('civitai_api')
+
+app = FastAPI(title="File Management API")
+
+# --- Pydantic Models for Request Validation ---
+class ProcessRequest(BaseModel):
+    urls: str # Accepts newline separated string to match original logic, or you could change to List[str]
+    hf_path: Optional[str] = "bigbossmonster/output"
+    selection: str = "default"  # options: default, wget, SmartDl, aria2, curl
+    hf_check: bool = False # If True, skip upload
+    upload_type: str = "model" # options: model, dataset
+    hf_token: Optional[str] = None # Optional override for token
+
+# --- Core Logic Functions (Refactored) ---
+
+def wipe_folder_logic():
+    try:
+        if os.path.exists(DIR):
+            shutil.rmtree(DIR)
+            os.makedirs(DIR, exist_ok=True)
+            return f"All files in '{DIR}' have been wiped successfully."
+        else:
+            return f"The folder '{DIR}' does not exist."
+    except Exception as e:
+        return f"Error wiping folder: {str(e)}"
+
+def extract_file_name(url):
+    parsed_url = urlparse(url)
+    path_segments = parsed_url.path.split('/')
+    return path_segments[-1]
+
+def download_file(url, save_path, selection):
+    if selection == "default":
+        try:
+            with requests.get(url, stream=True) as response:
+                response.raise_for_status()
+                with open(save_path, 'wb') as file:
+                    for chunk in response.iter_content(chunk_size=8192):
+                        file.write(chunk)
+            print(f"File downloaded to {save_path}")
+        except requests.RequestException as e:
+            print(f"Failed to download file: {e}")
+    elif selection in ["wget", "aria2", "curl"]:
+        try:
+            os.makedirs(DIR, exist_ok=True)
+            command = []
+            if selection == "wget":
+                command = ["wget", url, "-P", DIR]
+            elif selection == "aria2":
+                command = ["aria2c", "-x16", "-s16", "-d", DIR, url]
+            elif selection == "curl":
+                command = ["curl", "-o", save_path, url]
+            
+            subprocess.run(command, check=True)
+            print(f"File downloaded to {DIR} using {selection}")
+        except subprocess.CalledProcessError as e:
+            print(f"Failed to download file using {selection}: {e}")
+    else: # SmartDL
+        try:
+            obj = SmartDL(url, save_path)
+            obj.start()
+            print(f"File downloaded successfully to: {obj.get_dest()} using SmartDL")
+        except Exception as e:
+            print(f"Failed to download file using SmartDL: {e}")
+
+def extract_google_drive_id(url):
+    drive_id_pattern = re.compile(
+        r'(?:drive.google.com/.*?id=|drive.google.com/file/d/|drive.google.com/open\?id=|drive.google.com/uc\?id=|drive.google.com/drive/folders/)([a-zA-Z0-9_-]+)'
+    )
+    match = drive_id_pattern.search(url)
+    return match.group(1) if match else None
+
+def download_google_drive(url):
+    file_id = extract_google_drive_id(url)
+    if file_id:
+        download_url = f"https://drive.google.com/uc?id={file_id}"
+        print(f"Downloading from: {download_url}")
+        try:
+            # We must be inside the directory for gdown to work as intended with --folder sometimes
+            cwd = os.getcwd()
+            os.chdir(DIR)
+            if "folders" in url:
+                subprocess.run(["gdown", "--remaining-ok", "--folder" , url], check=True)
+            else:
+                subprocess.run(["gdown", download_url], check=True)
+            os.chdir(cwd) # Switch back
+            return f"Downloaded: {download_url}"
+        except subprocess.CalledProcessError as e:
+            os.chdir(cwd) # Switch back even on error
+            return f"Failed to download: {e}"
+        except Exception as e:
+             return f"Error: {e}"
+    else:
+        return f"Invalid Google Drive URL: {url}"
+
+def download_hugging_face(url, selection):
+    file_name = extract_file_name(url)
+    save_path = os.path.join(DIR, file_name)
+    download_file(url, save_path, selection)
+    return f"Downloaded: {url}"
+
+def download_katfile(url, selection):
+    apiKey = "699996yph6h88a7rc6c1g8"
+    try:
+        parts = url.split('/')
+        domain = parts[2]
+        filecode = parts[3]
+        cloneurl = f"https://{domain}/api/file/clone?key={apiKey}&file_code={filecode}"
+        
+        response = requests.get(cloneurl)
+        if response.status_code == 200:
+            json_data = response.json()
+            download_url = json_data.get('result', {}).get('url')
+            if not download_url: return "Failed to get Katfile URL"
+
+            parts_final = download_url.split('/')
+            filecodex = parts_final[3]
+            final_url = f"https://{domain}/api/file/direct_link?key={apiKey}&file_code={filecodex}"
+            
+            response = requests.get(final_url)
+            if response.status_code == 200:
+                json_data = response.json()
+                real_dl_url = json_data.get('result', {}).get('url')
+                save_path = os.path.join(DIR, extract_file_name(real_dl_url))
+                download_file(real_dl_url, save_path, selection)
+                return f"Downloaded: {real_dl_url}"
+    except Exception as e:
+        print(f"Error downloading from Katfile: {str(e)}")
+    return "Error downloading Katfile"
+
+def download_civitai(url, selection):
+    civit_key = get_civit_key()
+    parsed_url = urlparse(url)
+    path_components = parsed_url.path.split("/")
+    modelVersionId = path_components[-1]
+    save_path = os.path.join(DIR, f"{modelVersionId}.safecheck") # Temp name if unknown, usually CivitAI needs content-disposition
+    
+    # CivitAI specific logic
+    if selection == "curl":
+          command = [
+              "curl", "-L",
+              "-H", f"Authorization: Bearer {civit_key}",
+              "-o", save_path, # Added output path for curl
+              f"{url}"
+          ]
+          subprocess.run(command)
+    else:
+        # Construct API URL
+        api_url = f"https://civitai.com/api/download/models/{modelVersionId}?token={civit_key}"
+        # For standard downloads we generally want the real filename, 
+        # but sticking to your logic:
+        download_file(api_url, os.path.join(DIR, "civitai_model"), selection) 
+    
+    return f"Downloaded: {url}"
+
+def handle_download(url, selection):
+    try:
+        if "drive.google.com" in url:
+            return download_google_drive(url)
+        elif "huggingface.co" in url:
+            return download_hugging_face(url, selection)
+        elif "civitai.com" in url:
+            return download_civitai(url, selection)
+        elif "katfile.com" in url:
+            return download_katfile(url, selection)
+        else:
+            save_path = os.path.join(DIR, extract_file_name(url))
+            download_file(url, save_path, selection)
+            return f"Downloaded: {url}"
+    except Exception as e:
+        return f"Error occurred: {str(e)}"
+
+def upload_to_hf(hf_path, upload_type, token_val):
+    repo_id = hf_path if hf_path else "bigbossmonster/output"
+    commit_message = "Upload folder with Python script"
+
+    if os.path.exists(DIR):
+        api = HfApi()
+        try:
+            api.upload_folder(
+                folder_path=DIR,
+                repo_id=repo_id,
+                commit_message=commit_message,
+                token=token_val,
+                repo_type=upload_type,
+            )
+            return f"Upload success! Folder is available at https://huggingface.co/{repo_id}"
+        except Exception as e:
+            return f"Failed to upload folder. Error: {e}"
+    else:
+        return f"The folder '{DIR}' does not exist."
+
+# --- API Endpoints ---
+
+@app.get("/")
+def home():
+    return {"message": "File Management API is running. Go to /docs for Swagger UI."}
+
+@app.post("/process")
+def process_batch(request: ProcessRequest):
+    """
+    Main endpoint to download files and optionally upload to HF.
+    """
+    # 1. Process Downloads
+    download_results = []
+    urls_list = request.urls.splitlines()
+    
+    for url in urls_list:
+        url = url.strip()
+        if url:
+            result = handle_download(url, request.selection)
+            download_results.append(result)
+    
+    download_msg = "\n".join(download_results)
+    
+    # 2. Process Upload (if checked)
+    upload_msg = "Do Not upload to HF"
+    if not request.hf_check:
+        # Determine token: Request > Env Var
+        active_token = request.hf_token if request.hf_token else get_token()
+        upload_msg = upload_to_hf(request.hf_path, request.upload_type, active_token)
+    
+    return {
+        "download_status": download_msg,
+        "upload_status": upload_msg
+    }
+
+@app.post("/wipe")
+def wipe_folder():
+    """Wipes the download directory."""
+    return {"status": wipe_folder_logic()}
+
+@app.get("/files")
+def list_files_endpoint():
+    """Lists files in the download directory."""
+    try:
+        if not os.path.exists(DIR):
+            return {"files": [], "message": "Directory does not exist."}
+        
+        files = os.listdir(DIR)
+        return {"files": files, "count": len(files)}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)