ADD GDOWN

app.py

@@ -3,18 +3,31 @@ import requests
 import os
 import subprocess
 from typing import Union
-from datetime import datetime
 from pth2onnx import convert_pth_to_onnx
-from urllib.parse import urlparse
+import re
+import time
+from urllib.parse import urlparse, unquote
+import gdown
+import sys
 from typing import Optional
 
-# 新增日志开关
+# def format_bytes(size: int) -> str:
+#     """将字节数格式化为更易读的单位 (KB, MB, GB)"""
+#     if size is None:
+#         return "未知大小"
+#     power = 1024
+#     n = 0
+#     power_labels = {0: '', 1: 'K', 2: 'M', 3: 'G', 4: 'T'}
+#     while size > power and n < len(power_labels) -1 :
+#         size /= power
+#         n += 1
+#     return f"{size:.2f} {power_labels[n]}B"
+# 日志开关
 log_to_terminal = True
-
-# 新增全局任务计数器
 task_counter = 0
+download_cache = {}  # 格式: {url: 文件路径}
 
-# 新增日志函数
+# 日志函数
 def print_log(task_id, stage, status):
     if log_to_terminal:
         print(f"任务{task_id}: [{status}] {stage}")
@@ -31,28 +44,90 @@ def download_file(url: str, save_path: str):
     with open(save_path, 'wb') as f:
         f.write(response.content)
 
-def download_file2folder(url: str, folder: str, filesize_max: int, filesize_min: int) -> Optional[str]:
+
+
+def download_gdrive_file(
+    url: str, 
+    folder: str, 
+    filesize_max: int, 
+    filesize_min: int = 0
+) -> Optional[str]:
     """
-    从URL下载文件到指定文件夹，并进行文件大小检查
-    
+    从 Google Drive 链接获取文件信息，检查大小后下载文件。
+
     Args:
-        url: 要下载的文件URL
-        folder: 保存文件的目标文件夹路径
-        filesize_max: 最大允许文件大小(字节)，超过此值将中断下载并删除文件
-        filesize_min: 最小允许文件大小(字节)，小于此值将删除文件并返回None
-    
+        url (str): Google Drive 的分享链接。
+        folder (str): 用于保存文件的目标文件夹路径。
+        filesize_max (int): 允许下载的最大文件大小（单位：字节）。
+        filesize_min (int): 允许下载的最小文件大小（单位：字节），默认为 0。
+
     Returns:
-        成功下载的文件名，如果下载失败或文件大小不符合要求则返回None
+        Optional[str]: 如果下载成功，返回文件的完整路径；否则返回 None。
     """
-    # 解析URL获取文件名
-    parsed_url = urlparse(url)
-    filename = os.path.basename(parsed_url.path)
-    if not filename:
-        return None  # 无法从URL获取文件名
+    print(f"--- 开始处理链接: {url} ---")
     
-    # 确保目标文件夹存在
-    os.makedirs(folder, exist_ok=True)
-    save_path = os.path.join(folder, filename)
+    # 1. 获取文件信息
+    try:
+        info = gdown.get_file_info(url)
+        if not info:
+            print(f"错误：无法从链接获取文件信息。请检查链接是否有效且文件是否已公开分享。", file=sys.stderr)
+            return None
+            
+        filename = info.get('name')
+        filesize = info.get('size')
+
+        if not filename or filesize is None:
+            print(f"错误：无法获取完整的文件名或文件大小。元数据: {info}", file=sys.stderr)
+            return None
+
+    except Exception as e:
+        print(f"错误：获取文件信息时发生异常: {e}", file=sys.stderr)
+        return None
+
+    # 2. 验证文件大小
+    if filesize < filesize_min or filesize > filesize_max:
+        return None
+    
+
+    # 3. 准备下载路径并创建文件夹
+    try:
+        os.makedirs(folder, exist_ok=True)
+        output_path = os.path.join(folder, filename)
+    except OSError as e:
+        print(f"错误：创建文件夹 '{folder}' 失败: {e}", file=sys.stderr)
+        return None
+
+    # 4. 下载文件
+    try:
+        print(f"开始下载google drive文件到: {output_path}")
+        downloaded_path = gdown.download(url, output_path, quiet=True)
+        if downloaded_path and os.path.exists(downloaded_path):
+            print(f"下载成功！文件已保存至: {downloaded_path}")
+            return downloaded_path
+        else:
+            print("错误：gdown 下载过程未返回有效路径或文件下载后不存在。", file=sys.stderr)
+            return None
+            
+    except Exception as e:
+        print(f"错误：下载过程中发生异常: {e}", file=sys.stderr)
+        # 如果下载中断，清理可能不完整的文件
+        if os.path.exists(output_path):
+            os.remove(output_path)
+        return None
+
+
+
+def download_file2folder(url: str, folder: str, filesize_max: int, filesize_min: int) -> Optional[str]:
+    global download_cache
+    
+    # 检查缓存是否存在且文件有效
+    if url in download_cache:
+        cached_path = download_cache[url]
+        if os.path.exists(cached_path) and os.path.getsize(cached_path) >= filesize_min:
+            print(f"    使用缓存文件: {cached_path}")
+            return cached_path
+    
+    save_path = None  # 初始化save_path为None，避免except块中引用错误
     
     try:
         # 发送HTTP请求，流式下载
@@ -64,6 +139,32 @@ def download_file2folder(url: str, folder: str, filesize_max: int, filesize_min:
             if total_size > filesize_max:
                 return None  # 文件大小超过最大值，不下载
             
+            # 提取文件名
+            filename = None
+            content_disposition = response.headers.get('Content-Disposition')
+            if content_disposition:
+                # 使用正则表达式提取filename
+                match = re.search(r'filename="?([^"]+)"?', content_disposition)
+                if match:
+                    filename = match.group(1)
+                    # 处理可能的URL编码文件名
+                    filename = unquote(filename)
+            
+            # 如果响应头没有，从URL解析
+            if not filename:
+                parsed_url = urlparse(url)
+                filename = os.path.basename(parsed_url.path)
+                # 处理URL编码的文件名
+                filename = unquote(filename)
+            
+            # 如果仍然没有文件名，生成默认文件名
+            if not filename:
+                filename = f"download_{int(time.time())}.bin"
+            
+            # 确保目标文件夹存在
+            os.makedirs(folder, exist_ok=True)
+            save_path = os.path.join(folder, filename)
+            
             downloaded_size = 0
             with open(save_path, 'wb') as file:
                 for chunk in response.iter_content(chunk_size=8192):
@@ -81,11 +182,14 @@ def download_file2folder(url: str, folder: str, filesize_max: int, filesize_min:
             os.remove(save_path)
             return None
         
-        return filename
+        # 下载成功后更新缓存
+        if save_path:
+            download_cache[url] = save_path
+        return save_path
     
     except Exception as e:
         # 发生异常时清理文件
-        if os.path.exists(save_path):
+        if save_path and os.path.exists(save_path):
             os.remove(save_path)
         return None
 
@@ -154,14 +258,15 @@ async def _process_model(model_input: Union[str, gr.File], tilesize: int, output
 
 with gr.Blocks() as demo:
     gr.Markdown("# 模型转换工具")
+    model_type_opt = ['从链接下载', '直接上传文件']
     with gr.Row():
         with gr.Column():
-            input_type = gr.Radio(['模型链接', '上传模型文件'], label='输入类型')
+            input_type = gr.Radio(model_type_opt, label='模型文件来源')
             url_input = gr.Textbox(label='模型链接')
-            file_input = gr.File(label='上传模型文件', visible=False)
+            file_input = gr.File(label='模型文件', visible=False)
 
             def show_input(input_type):
-                if input_type == '模型链接':
+                if input_type == model_type_opt[0]:
                     return gr.update(visible=True), gr.update(visible=False)
                 else:
                     return gr.update(visible=False), gr.update(visible=True)
@@ -186,30 +291,43 @@ with gr.Blocks() as demo:
         os.makedirs(output_dir, exist_ok=True)
         log=""
 
-        if input_type == '模型链接' and url_input:
-            # 新增：下载模型文件到 output_dir
+        if input_type == model_type_opt[0] and url_input:
             log = f'正在下载模型文件: {url_input}\n'
             print_log(task_counter, f'正在下载模型文件: {url_input}', '开始')
             yield None, None, log
-        
-            filename = download_file2folder(
-                url=url_input,
-                folder=output_dir,
-                filesize_max=200*1024*1024,  # 200MB
-                filesize_min=1024     # 1KB
-            )
-            model_input = os.path.join(output_dir, filename)
+
+            if url_input.startswith("https://drive.google.com/file/"):
+                model_input = download_gdrive_file(
+                    url=url_input,
+                    folder=output_dir,
+                    filesize_max=200*1024*1024,  # 200MB
+                    filesize_min=1024     # 1KB
+                )
+            else:
+                model_input = download_file2folder(
+                    url=url_input,
+                    folder=output_dir,
+                    filesize_max=200*1024*1024,  # 200MB
+                    filesize_min=1024     # 1KB
+                )
+
+
+            if not model_input:
+                log += f'\n模型文件下载失败\n'
+                print_log(task_counter, f'模型文件载', '失败')
+                yield None, None, log
+                return
+            
             log += f'\n��型文件已下载到: {model_input}\n'
             print_log(task_counter, f'模型文件已下载到: {model_input}', '完成')
             yield None, None, log
-        elif input_type == '上传模型文件' and file_input:
+        elif input_type == model_type_opt[1] and file_input:
             model_input = file_input
         else:
             # 改为通过yield返回错误日志
             log = '\n请选择输入类型并提供有效的输入！'
             yield None, None, log
             return
-
         
         onnx_path = None
         mnn_path = None
@@ -235,11 +353,13 @@ with gr.Blocks() as demo:
     examples_column = gr.Column(visible=True)
     with examples_column:
         examples = [
-            ["模型链接", "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth"],
-            ["模型链接", "https://github.com/Phhofm/models/releases/download/4xNomos8kSC/4xNomos8kSC.pth"],
-            ["模型链接", "https://github.com/Phhofm/models/releases/download/1xDeJPG/1xDeJPG_SRFormer_light.pth"],
-            ["模型链接", "https://objectstorage.us-phoenix-1.oraclecloud.com/n/ax6ygfvpvzka/b/open-modeldb-files/o/4x-WTP-ColorDS.pth"],
-            ["模型链接", "https://github.com/Kim2091/Kim2091-Models/releases/download/2x-AnimeSharpV4/2x-AnimeSharpV4_RCAN_fp16_op17.onnx"]
+            [model_type_opt[0], "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth"],
+            [model_type_opt[0], "https://github.com/Phhofm/models/releases/download/4xNomos8kSC/4xNomos8kSC.pth"],
+            [model_type_opt[0], "https://github.com/Phhofm/models/releases/download/1xDeJPG/1xDeJPG_SRFormer_light.pth"],
+            [model_type_opt[0], "https://objectstorage.us-phoenix-1.oraclecloud.com/n/ax6ygfvpvzka/b/open-modeldb-files/o/4x-WTP-ColorDS.pth"],
+            [model_type_opt[0], "https://github.com/Kim2091/Kim2091-Models/releases/download/2x-AnimeSharpV4/2x-AnimeSharpV4_RCAN_fp16_op17.onnx"],
+            [model_type_opt[0], "https://drive.google.com/uc?export=download&confirm=1&id=1PeqL1ikJbBJbVzvlqvtb4d7QdSW7BzrQ"],
+            [model_type_opt[0], "https://drive.google.com/file/d/1maYmC5yyzWCC42X5O0HeDuepsLFh7AV4/view?usp=drive_link"],
         ]
         example_input = gr.Examples(examples=examples, inputs=[input_type, url_input], label='示例模型链接')
     

pth2onnx.py

@@ -73,7 +73,7 @@ def convert_pth_to_onnx(pth_path: str, onnx_path: str=None, channel:int=0, tiles
             print(f'文件名 {filename} 包含匹配模式。')
         else:
             base_path = f"{base_path}-x{scale}"
-        print("final use_fp16", str(use_fp16) )
+        # print("final use_fp16", str(use_fp16) )
         onnx_path = base_path + ("-Grayscale" if channel==1 else "") + ("-fp16.onnx" if use_fp16 else ".onnx")
         
     # 处理相对路径情况

requirements.txt

@@ -5,3 +5,4 @@ onnx
 onnxsim
 mnn
 gradio
+gdown