Upload app.py

app.py

@@ -40,6 +40,19 @@ pipe = DAIPipeline(
     t_start=0,
 ).to(device)
 
+def resize_image(image, max_size):
+    """Resize the image so that the maximum side is max_size."""
+    width, height = image.size
+    if max(width, height) > max_size:
+        if width > height:
+            new_width = max_size
+            new_height = int(height * (max_size / width))
+        else:
+            new_height = max_size
+            new_width = int(width * (max_size / height))
+        image = image.resize((new_width, new_height), Image.LANCZOS)
+    return image
+
 @spaces.GPU
 def process_image(input_image, resolution_choice):
     # 将 Gradio 输入转换为 PIL 图像
@@ -47,6 +60,7 @@ def process_image(input_image, resolution_choice):
 
     # 根据用户选择设置处理分辨率
     if resolution_choice == "768":
+        input_image = resize_image(input_image, 768)
         processing_resolution = None
     else:
         processing_resolution = 0  # 使用原始分辨率
@@ -64,7 +78,7 @@ def process_image(input_image, resolution_choice):
     processed_frame = (processed_frame[0] * 255).astype(np.uint8)
     processed_frame = Image.fromarray(processed_frame)
 
-    return processed_frame
+    return input_image, processed_frame  # 返回调整后的输入图片和处理后的图片
 
 # 创建 Gradio 界面
 def create_gradio_interface():
@@ -108,7 +122,7 @@ def create_gradio_interface():
         submit_btn.click(
             fn=process_image,
             inputs=[input_image, resolution_choice],  # 输入组件列表
-            outputs=output_image,
+            outputs=[input_image, output_image],  # 输出组件列表
         )
 
     return demo