update

app.py

@@ -10,19 +10,8 @@ from urllib.parse import urlparse, unquote
 import gdown
 import sys
 from typing import Optional
+import datetime
 
-# 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: 文件路径}
@@ -44,8 +33,6 @@ def download_file(url: str, save_path: str):
     with open(save_path, 'wb') as f:
         f.write(response.content)
 
-
-
 def download_gdrive_file(
     url: str, 
     folder: str, 
@@ -65,57 +52,53 @@ def download_gdrive_file(
         Optional[str]: 如果下载成功，返回文件的完整路径；否则返回 None。
     """
     print(f"--- 开始处理链接: {url} ---")
-    
-    # 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
+        print(f"开始下载google drive文件到: {folder}")
+        # 保存当前工作目录
+        original_dir = os.getcwd()
+        # 切换到目标文件夹
+        os.chdir(folder)
+        # 下载文件，不指定output_path参数
+        downloaded_filename = gdown.download(url, quiet=True, fuzzy=True)
+        # 恢复原始工作目录
+        os.chdir(original_dir)
+        
+        # 检查下载结果和文件信息
+        if downloaded_filename:
+            downloaded_path = os.path.join(folder, downloaded_filename)
+            if os.path.exists(downloaded_path):
+                # 获取文件大小
+                file_size = os.path.getsize(downloaded_path)
+                if file_size< filesize_max and file_size> filesize_min:
+                    print(f"下载成功！文件已保存至: {downloaded_path}")
+                    return downloaded_path
+                else:
+                    print(f"文件大小超出范围: {file_size} bytes")
+                    return None
+            else:
+                print(f"错误：下载的文件 '{downloaded_path}' 不存在。", file=sys.stderr)
+                return None
         else:
-            print("错误：gdown 下载过程未返回有效路径或文件下载后不存在。", file=sys.stderr)
+            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)
+        print(f"下载过程中发生错误: {e}", file=sys.stderr)
         return None
-
-
+    finally:
+        # 确保工作目录恢复，即使发生异常
+        if 'original_dir' in locals():
+            os.chdir(original_dir)
+            
 
 def download_file2folder(url: str, folder: str, filesize_max: int, filesize_min: int) -> Optional[str]:
     global download_cache
@@ -287,7 +270,9 @@ with gr.Blocks() as demo:
     async def process_model(input_type, url_input, file_input, tilesize, fp16, try_run):
         global task_counter
         task_counter += 1
-        output_dir = os.path.join(os.getcwd(), f"output_{task_counter}")
+
+        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+        output_dir = os.path.join(os.getcwd(), f"output_{task_counter}_{timestamp}")
         os.makedirs(output_dir, exist_ok=True)
         log=""
 
@@ -296,7 +281,7 @@ with gr.Blocks() as demo:
             print_log(task_counter, f'正在下载模型文件: {url_input}', '开始')
             yield None, None, log
 
-            if url_input.startswith("https://drive.google.com/file/"):
+            if url_input.startswith("https://drive.google.com/"):
                 model_input = download_gdrive_file(
                     url=url_input,
                     folder=output_dir,

pth2onnx.py

@@ -70,7 +70,7 @@ def convert_pth_to_onnx(pth_path: str, onnx_path: str=None, channel:int=0, tiles
         filename = os.path.basename(pth_path).upper()
         pattern = f'(^|[_-])({scale}X|X{scale})([_-]|$)'
         if re.search(pattern, filename):
-            print(f'文件名 {filename} 包含匹配模式。')
+            print(f'文件名 {filename} 包含倍率信息。')
         else:
             base_path = f"{base_path}-x{scale}"
         # print("final use_fp16", str(use_fp16) )