update app.py

app.py

@@ -31,7 +31,20 @@ def start_process(input1, input2, shape0_str, shape1_str, input_suffix=".pth"):
     log = "转换过程非常慢，请耐心等待。显示文件列表不代表转换完成。如果未发生错误，转换结束会显示”任务完成“\n"
     yield [], log
     if input2 == None or input2.strip() == "":
-        input2 = os.path.splitext(os.path.basename(input1))[0]
+        split_input = os.path.splitext(os.path.basename(input1))
+        if len(split_input) > 1:
+            suffix = split_input[1].split('?')[0].lower()
+            if  suffix not in [".pth" , ".safetensors" , ".ckpt"]:
+                print_log(task_id, input2,  "不支持此文件的格式 suffix="+suffix, "错误")
+                log += f"不支持此文件的格式\n"
+                return [] , log
+        print_log(task_id, input2, "检查文件名", "开始")
+        log += f"检查文件名…\n"
+        yield [], log
+        if input2 == None or input2.strip() == "":
+            input2 = str(task_id)
+
+        input1 = split_input[0]
         if input2 == "":
             input2 = str(task_id)
         log += f"未提供文件名，使用{input2}\n"
@@ -245,7 +258,7 @@ with gr.Blocks() as demo:
             with gr.Row():
                 input1 = gr.Textbox(label="粘贴地址")
                 # 新增文件上传组件
-                input1_file = gr.File(label="上传文件")
+                input1_file = gr.File(label="上传文件", file_types=[".pth", ".safetensors", ".ckpt"])
             input2 = gr.Textbox(label="自定义文件名")
             # 修改为字符串输入控件
             shape0_str = gr.Textbox(label="shape0 (逗号分隔的整数)", value="1,3,128,128")
@@ -275,13 +288,13 @@ with gr.Blocks() as demo:
 
     # 添加范例
     examples = [
-        ["https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth", None, "", "1,3,128,128", "0,0,0,0"],
-        ["https://github.com/Phhofm/models/releases/download/4xNomos8kSC/4xNomos8kSC.pth", None, "", "1,3,128,128", "0,0,0,0"],
-        ["https://github.com/Phhofm/models/releases/download/1xDeJPG/1xDeJPG_SRFormer_light.pth", None, "", "1,3,128,128", "0,0,0,0"]
+        ["https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth", "", "1,3,128,128", "0,0,0,0"],
+        ["https://github.com/Phhofm/models/releases/download/4xNomos8kSC/4xNomos8kSC.pth", "", "1,3,128,128", "0,0,0,0"],
+        ["https://github.com/Phhofm/models/releases/download/1xDeJPG/1xDeJPG_SRFormer_light.pth", "", "1,3,128,128", "0,0,0,0"]
     ]
     gr.Examples(
         examples=examples,
-        inputs=[input1, input1_file, input2, shape0_str, shape1_str],
+        inputs=[input1, input2, shape0_str, shape1_str],
         outputs=[output, log_textbox],
         fn=start_process
     )