update code

app.py

@@ -4,6 +4,7 @@ import sys
 import warnings
 import logging
 import spaces
+import difflib
 
 # Configure logging settings
 logging.basicConfig(
@@ -635,7 +636,9 @@ def process_image(input_image_dir, source_type, is_style, save_dir):
     """ 🎯 处理 input_image，根据是否是示例图片执行不同逻辑 """
     process_img_input_dir = os.path.join(save_dir, 'input_image')
     process_img_save_dir = os.path.join(save_dir, 'processed_img')
-    image_name_true = os.path.basename(input_image_dir)
+    base_name = os.path.basename(input_image_dir)           # abc123.jpg
+    name_without_ext = os.path.splitext(base_name)[0]        # abc123
+    image_name_true  = name_without_ext + ".png" 
     os.makedirs(process_img_save_dir, exist_ok=True)
     os.makedirs(process_img_input_dir, exist_ok=True)
     if source_type == "example":
@@ -645,8 +648,15 @@ def process_image(input_image_dir, source_type, is_style, save_dir):
         # input_process_model.inference(input_image, process_img_save_dir)
         shutil.copy(input_image_dir, process_img_input_dir)
         input_process_model.inference(process_img_input_dir, process_img_save_dir, is_img=True, is_video=False)
-        imge_dir = os.path.join(save_dir, 'processed_img/dataset/images512x512/input_image', image_name_true)
+
+        files = os.listdir(os.path.join(process_img_save_dir, 'dataset/images512x512/input_image'))
+        image_files = [f for f in files if f.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.webp'))]
+        # 使用 difflib 查找相似文件名
+        matches = difflib.get_close_matches(image_name_true, image_files, n=1, cutoff=0.1)
+        closest_match = matches[0]
+        imge_dir = os.path.join(process_img_save_dir, 'processed_img/dataset/images512x512/input_image', closest_match)
         image = Image.open(imge_dir)
+        image_name_true = closest_match
         return image, source_type, image_name_true  # 这里替换成 处理用户上传图片的逻辑
 
 @spaces.GPU(duration=100)