Update app.py

app.py

@@ -1,224 +1,125 @@
-import requests
-import os
 import gradio as gr
-from huggingface_hub import update_repo_visibility, whoami, upload_folder, create_repo, upload_file
-from slugify import slugify
-import re
-import uuid
-from typing import Optional
-import json
-
-# List of trusted uploaders
-TRUSTED_UPLOADERS = [
-    "KappaNeuro", "CiroN2022", "multimodalart", "Norod78", "joachimsallstrom", 
-    "blink7630", "e-n-v-y", "DoctorDiffusion", "RalFinger", "artificialguybr", "nevreal"
-]
-
-# Fetch JSON data from CivitAI
-def get_json_data(url):
-    url_split = url.split('/')
-    api_url = f"https://civitai.com/api/v1/models/{url_split[4]}"
+import os
+import time
+import urllib.request
+from pathlib import Path
+from urllib.parse import urlparse, parse_qs, unquote
+
+CHUNK_SIZE = 1638400
+TOKEN_FILE = Path.home() / '.civitai' / 'config'
+USER_AGENT = 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/58.0.3029.110 Safari/537.3'
+
+def get_token():
     try:
-        response = requests.get(api_url)
-        response.raise_for_status()
-        return response.json()
-    except requests.exceptions.RequestException as e:
-        print(f"Error fetching JSON data: {e}")
+        with open(TOKEN_FILE, 'r') as file:
+            token = file.read()
+            return token
+    except Exception:
         return None
 
-# Fetch prompt data from images using image IDs
-def get_prompts_from_image(image_id):
-    url = f'https://civitai.com/api/trpc/image.getGenerationData?input={{"json":{{"id":{image_id}}}}}'
-    response = requests.get(url)
-    prompt = ""
-    negative_prompt = ""
-    if response.status_code == 200:
-        data = response.json()
-        result = data['result']['data']['json']
-        if "prompt" in result['meta']:
-            prompt = result['meta']['prompt']
-        if "negativePrompt" in result['meta']:
-            negative_prompt = result["meta"]["negativePrompt"]
-    return prompt, negative_prompt
-
-# Extract model information from JSON data
-def extract_info(json_data):
-    if json_data["type"] == "LORA":
-        for model_version in json_data["modelVersions"]:
-            if model_version["baseModel"] in [
-                "SDXL 1.0", "SDXL 0.9", "SD 1.5", "SD 1.4", "SD 2.1", "SD 2.0", 
-                "SD 2.0 768", "SD 2.1 768", "SD 3", "Flux.1 D", "Flux.1 S"
-            ]:
-                for file in model_version["files"]:
-                    if "primary" in file:
-                        urls_to_download = [{
-                            "url": file["downloadUrl"], 
-                            "filename": file["name"], 
-                            "type": "weightName"
-                        }]
-                        
-                        for image in model_version["images"]:
-                            image_id = image["url"].split("/")[-1].split(".")[0]
-                            prompt, negative_prompt = get_prompts_from_image(image_id)
-                            urls_to_download.append({
-                                "url": image["url"],
-                                "filename": os.path.basename(image["url"]),
-                                "type": "imageName",
-                                "prompt": prompt,
-                                "negative_prompt": negative_prompt
-                            })
-                        model_mapping = {
-                            "SDXL 1.0": "stabilityai/stable-diffusion-xl-base-1.0",
-                            "SDXL 0.9": "stabilityai/stable-diffusion-xl-base-1.0",
-                            "SD 1.5": "runwayml/stable-diffusion-v1-5",
-                            "SD 1.4": "CompVis/stable-diffusion-v1-4",
-                            "SD 2.1": "stabilityai/stable-diffusion-2-1-base",
-                            "SD 2.0": "stabilityai/stable-diffusion-2-base",
-                            "SD 2.1 768": "stabilityai/stable-diffusion-2-1",
-                            "SD 2.0 768": "stabilityai/stable-diffusion-2",
-                            "SD 3": "stabilityai/stable-diffusion-3-medium-diffusers",
-                            "Flux.1 D": "black-forest-labs/FLUX.1-dev",
-                            "Flux.1 S": "black-forest-labs/FLUX.1-schnell"
-                        }
-                        base_model = model_mapping[model_version["baseModel"]]
-                        info = {
-                            "urls_to_download": urls_to_download,
-                            "id": model_version["id"],
-                            "baseModel": base_model,
-                            "modelId": model_version.get("modelId", ""),
-                            "name": json_data["name"],
-                            "description": json_data["description"],
-                            "trainedWords": model_version.get("trainedWords", []),
-                            "creator": json_data["creator"]["username"],
-                            "tags": json_data["tags"],
-                            "allowNoCredit": json_data["allowNoCredit"],
-                            "allowCommercialUse": json_data["allowCommercialUse"],
-                            "allowDerivatives": json_data["allowDerivatives"],
-                            "allowDifferentLicense": json_data["allowDifferentLicense"]
-                        }
-                        return info
-    return None
-
-# Function to download files
-def download_files(info, folder="."):
-    downloaded_files = {
-        "imageName": [],
-        "imagePrompt": [],
-        "imageNegativePrompt": [],
-        "weightName": []
+def store_token(token: str):
+    TOKEN_FILE.parent.mkdir(parents=True, exist_ok=True)
+    with open(TOKEN_FILE, 'w') as file:
+        file.write(token)
+
+def download_file(url: str, output_path: str, token: str):
+    headers = {
+        'Authorization': f'Bearer {token}',
+        'User-Agent': USER_AGENT,
     }
-    for item in info["urls_to_download"]:
-        download_file(item["url"], item["filename"], folder)
-        downloaded_files[item["type"]].append(item["filename"])
-        if item["type"] == "imageName":
-            prompt_clean = re.sub(r'<.*?>', '', item["prompt"])
-            negative_prompt_clean = re.sub(r'<.*?>', '', item["negative_prompt"])
-            downloaded_files["imagePrompt"].append(prompt_clean)
-            downloaded_files["imageNegativePrompt"].append(negative_prompt_clean)
-    return downloaded_files
-
-# Function to download individual files
-def download_file(url, filename, folder="."):
-    headers = {}
-    try:
-        response = requests.get(url, headers=headers)
-        response.raise_for_status()
-    except requests.exceptions.HTTPError as e:
-        print(e)
-        if response.status_code == 401: 
-            headers['Authorization'] = f'Bearer {os.environ["CIVITAI_API"]}'
-            try:
-                response = requests.get(url, headers=headers)
-                response.raise_for_status()
-            except requests.exceptions.RequestException as e:
-                raise gr.Error(f"Error downloading file: {e}")
-        else:
-            raise gr.Error(f"Error downloading file: {e}")
-    except requests.exceptions.RequestException as e:
-        raise gr.Error(f"Error downloading file: {e}")
-
-    with open(f"{folder}/{filename}", 'wb') as f:
-        f.write(response.content)
-
-# Main function to process URL and extract necessary info
-def process_url(url, profile, do_download=True, folder="."):
-    json_data = get_json_data(url)
-    if json_data:
-        if check_nsfw(json_data, profile):
-            info = extract_info(json_data)
-            if info:
-                if do_download:
-                    downloaded_files = download_files(info, folder)
-                else:
-                    downloaded_files = []
-                return info, downloaded_files
-            else:
-                raise gr.Error("Only SDXL LoRAs are supported for now")
+
+    class NoRedirection(urllib.request.HTTPErrorProcessor):
+        def http_response(self, request, response):
+            return response
+        https_response = http_response
+
+    request = urllib.request.Request(url, headers=headers)
+    opener = urllib.request.build_opener(NoRedirection)
+    response = opener.open(request)
+
+    if response.status in [301, 302, 303, 307, 308]:
+        redirect_url = response.getheader('Location')
+        parsed_url = urlparse(redirect_url)
+        query_params = parse_qs(parsed_url.query)
+        content_disposition = query_params.get('response-content-disposition', [None])[0]
+
+        if content_disposition:
+            filename = unquote(content_disposition.split('filename=')[1].strip('"'))
         else:
-            raise gr.Error("This model has content tagged as unsafe by CivitAI")
+            return None, f"ERROR: Unable to determine filename"
+
+        response = urllib.request.urlopen(redirect_url)
+    elif response.status == 404:
+        return None, f"ERROR: File not found"
     else:
-        raise gr.Error("Something went wrong in fetching CivitAI API")
+        return None, f"ERROR: No redirect found, something went wrong"
 
-# Creating the web UI using Gradio
-with gr.Blocks() as demo:
-    gr.Markdown("# Upload your CivitAI LoRA to Hugging Face 🤗")
-    
-    # Login Button
-    gr.LoginButton()
-    
-    # URL input and submit button
-    submit_source_civit = gr.Textbox(
-        placeholder="https://civitai.com/models/144684/pixelartredmond-pixel-art-loras-for-sd-xl",
-        label="CivitAI model URL",
-        info="URL of the CivitAI LoRA"
-    )
-    submit_button_civit = gr.Button("Upload model to Hugging Face and submit")
+    total_size = response.getheader('Content-Length')
+    if total_size is not None:
+        total_size = int(total_size)
+
+    output_file = os.path.join(output_path, filename)
+
+    with open(output_file, 'wb') as f:
+        downloaded = 0
+        start_time = time.time()
+
+        while True:
+            chunk_start_time = time.time()
+            buffer = response.read(CHUNK_SIZE)
+            chunk_end_time = time.time()
+
+            if not buffer:
+                break
+
+            downloaded += len(buffer)
+            f.write(buffer)
+            chunk_time = chunk_end_time - chunk_start_time
+
+            if chunk_time > 0:
+                speed = len(buffer) / chunk_time / (1024 ** 2)  # Speed in MB/s
+
+            if total_size is not None:
+                progress = downloaded / total_size
+                gr.Progress.update(f'Downloading: {filename} [{progress*100:.2f}%] - {speed:.2f} MB/s')
 
-    # Output section for progress and result display
-    output = gr.Markdown(label="Output progress", visible=False)
+    end_time = time.time()
+    time_taken = end_time - start_time
+    hours, remainder = divmod(time_taken, 3600)
+    minutes, seconds = divmod(remainder, 60)
+
+    if hours > 0:
+        time_str = f'{int(hours)}h {int(minutes)}m {int(seconds)}s'
+    elif minutes > 0:
+        time_str = f'{int(minutes)}m {int(seconds)}s'
+    else:
+        time_str = f'{int(seconds)}s'
+
+    return output_file, f"Download completed. File saved as: {filename}. Downloaded in {time_str}"
+
+def run_downloader(url, output_path, token_input):
+    token = get_token() or token_input
+    if not token:
+        return None, "ERROR: Token not provided."
+
+    store_token(token)
+    file_path, message = download_file(url, output_path, token)
+    return file_path, message
+
+with gr.Blocks() as demo:
+    gr.Markdown("# CivitAI Downloader")
+    with gr.Row():
+        url_input = gr.Textbox(label="CivitAI Download URL")
+        output_path_input = gr.Textbox(label="Output Path", value="./")  # Default value set to current directory
+        token_input = gr.Textbox(label="CivitAI API Token (Optional)")
+    download_button = gr.Button("Download")
+    output_file = gr.File(label="Downloaded Model")
+    message_output = gr.Textbox(label="Status Message", interactive=False)
     
-    # Functionality to handle button click events
-    def upload_model_to_hf(url, profile, oauth_token):
-        if not profile.name:
-            return gr.Error("Are you sure you are logged in?")
-        
-        folder = str(uuid.uuid4())
-        os.makedirs(folder, exist_ok=False)
-        info, downloaded_files = process_url(url, profile, folder=folder)
-        slug_name = slugify(info["name"])
-        user_repo_id = f"{profile.username}/{slug_name}"
-        
-        # Create repository on Hugging Face
-        try:
-            create_repo(repo_id=user_repo_id, repo_type="model", space_sdk="diffusers", private=True, token=oauth_token)
-        except Exception as e:
-            raise gr.Error(f"Error creating Hugging Face repo: {e}")
-        
-        try:
-            upload_folder(
-                folder_path=folder,
-                repo_id=user_repo_id,
-                repo_type="model",
-                token=oauth_token
-            )
-        except Exception as e:
-            raise gr.Error(f"Error uploading to Hugging Face: {e}")
-        
-        update_repo_visibility(repo_id=user_repo_id, private=False, token=oauth_token)
-        
-        # Clean up downloaded files
-        for filename in downloaded_files:
-            os.remove(os.path.join(folder, filename))
-        os.rmdir(folder)
-        
-        return f"Model '{info['name']}' uploaded to {user_repo_id}"
-
-    # Define the submit button click action
-    submit_button_civit.click(
-        upload_model_to_hf, 
-        inputs=[submit_source_civit, gr.OAuthProfile(hub_token=gr.Token())], 
-        outputs=output
+    download_button.click(
+        run_downloader, 
+        inputs=[url_input, output_path_input, token_input],
+        outputs=[output_file, message_output]
     )
 
-# Launch the Gradio interface
 demo.launch()