Update app.py

app.py

@@ -9,6 +9,7 @@ from PIL import Image
 from typing import Iterable
 from gradio.themes import Soft
 from gradio.themes.utils import colors, fonts, sizes
+from datetime import datetime
 
 colors.orange_red = colors.Color(
     name="orange_red",
@@ -130,15 +131,10 @@ def update_dimensions_on_upload(image):
     return new_width, new_height
 
 @spaces.GPU
-def infer(
-    images,
-    prompt,
-    seed,
-    randomize_seed,
-    guidance_scale,
-    steps,
-    progress=gr.Progress(track_tqdm=True)
-):
+def run_gpu_inference(images, prompt, seed, randomize_seed, guidance_scale, steps):
+    """
+    负责在 ZeroGPU 上运行模型推理，运行完毕后返回结果并自动释放 GPU
+    """
     gc.collect()
     torch.cuda.empty_cache()
 
@@ -187,7 +183,8 @@ def infer(
             true_cfg_scale=guidance_scale,
         ).images[0]
 
-        return result_image, seed
+        # 返回：生成的结果图，最终使用的 seed，以及被用作底图的第一张原始图片
+        return result_image, seed, pil_images[0]
 
     except Exception as e:
         raise e
@@ -195,25 +192,37 @@ def infer(
         gc.collect()
         torch.cuda.empty_cache()
 
-@spaces.GPU
-def infer_example(images, prompt):
-    if not images:
-        return None, 0
-
-    if isinstance(images, str):
-        images_list = [images]
-    else:
-        images_list = images
-
-    result, seed = infer(
-        images=images_list,
-        prompt=prompt,
-        seed=0,
-        randomize_seed=True,
-        guidance_scale=1.0,
-        steps=4
+def process_and_concat(images, prompt, seed, randomize_seed, guidance_scale, steps, progress=gr.Progress(track_tqdm=True)):
+    """
+    负责调用 GPU 推理，等待 GPU 释放后，在 CPU 环境下完成上下拼接并生成规范命名的无损图片。
+    """
+    # 1. 执行 GPU 推理 (此函数运行完毕后，ZeroGPU 即会关闭/释放)
+    result_image, final_seed, orig_image = run_gpu_inference(
+        images, prompt, seed, randomize_seed, guidance_scale, steps
     )
-    return result, seed
+
+    # 2. 修改前后上下拼接 (在 CPU 上进行)
+    # 为了保证拼接整齐，将原图 resize 为与结果图完全一致的尺寸
+    orig_resized = orig_image.resize((result_image.width, result_image.height), Image.Resampling.LANCZOS)
+    
+    total_width = result_image.width
+    total_height = orig_resized.height + result_image.height
+    
+    concat_img = Image.new('RGB', (total_width, total_height))
+    # 上：原图
+    concat_img.paste(orig_resized, (0, 0))
+    # 下：修改后的图
+    concat_img.paste(result_image, (0, orig_resized.height))
+
+    # 3. 保存为规范命名的无损 PNG 格式
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    filename = f"[edit]_{timestamp}_seed{final_seed}.png"
+    filepath = os.path.join(os.getcwd(), filename)
+    concat_img.save(filepath, format="PNG")
+
+    # 返回拼接后的预览图，提供下载的文件路径，以及 seed
+    return concat_img, filepath, final_seed
+
 
 css = """
 #col-container {
@@ -232,7 +241,6 @@ with gr.Blocks() as demo:
             with gr.Column():
                 images = gr.Gallery(
                     label="Upload Images",
-                    #sources=["upload", "clipboard"],
                     type="filepath",
                     columns=2,
                     rows=1,
@@ -250,7 +258,9 @@ with gr.Blocks() as demo:
                 run_button = gr.Button("Edit Image", variant="primary")
 
             with gr.Column():
-                output_image = gr.Image(label="Output Image", interactive=False, format="png", height=395)
+                output_image = gr.Image(label="Output Image (Top: Before, Bottom: After)", interactive=False, format="png")
+                # 新增用于一键下载无损图的组件
+                output_file = gr.File(label="Download Lossless Merged Image")
 
                 with gr.Accordion("Advanced Settings", open=False, visible=False):
                     seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
@@ -258,28 +268,12 @@ with gr.Blocks() as demo:
                     guidance_scale = gr.Slider(label="Guidance Scale", minimum=1.0, maximum=10.0, step=0.1, value=1.0)
                     steps = gr.Slider(label="Inference Steps", minimum=1, maximum=50, step=1, value=4)
 
-        gr.Examples(
-            examples=[
-                [["examples/1.jpg"], "cinematic polaroid with soft grain subtle vignette gentle lighting white frame handwritten photographed 'Fire-Edit' preserving realistic texture and details."],
-                [["examples/2.jpg"], "Transform the image into a dotted cartoon style."],
-                [["examples/3.jpeg"], "Convert it to black and white."],
-                [["examples/4.jpg", "examples/5.jpg"], "Replace her glasses with the new glasses from image 1."],
-                [["examples/8.jpg", "examples/9.png"], "Replace the current clothing with the clothing from the reference image 2. Keep the person’s face, hairstyle, body pose, background, lighting, and camera angle unchanged. Ensure the new outfit fits naturally with realistic fabric texture, proper shadows, folds, and accurate proportions. Match the lighting, color tone, and overall style for a seamless and high-quality result."],
-                [["examples/10.jpg", "examples/11.png"], "Replace the current clothing with the clothing from the reference image 2. Keep the person’s face, hairstyle, body pose, background, lighting, and camera angle unchanged. Ensure the new outfit fits naturally with realistic fabric texture, proper shadows, folds, and accurate proportions. Match the lighting, color tone, and overall style for a seamless and high-quality result."],
-            ],
-            inputs=[images, prompt],
-            outputs=[output_image, seed],
-            fn=infer_example,
-            cache_examples=False,
-            label="Examples"
-        )
-        
         gr.Markdown("[*](https://huggingface.co/FireRedTeam/FireRed-Image-Edit-1.0)This is still an experimental Space for FireRed-Image-Edit-1.0.")
 
     run_button.click(
-        fn=infer,
+        fn=process_and_concat,
         inputs=[images, prompt, seed, randomize_seed, guidance_scale, steps],
-        outputs=[output_image, seed]
+        outputs=[output_image, output_file, seed]
     )
 
 if __name__ == "__main__":