added paste back as an option

app.py

@@ -46,7 +46,7 @@ pipe.scheduler = TCDScheduler.from_config(pipe.scheduler.config)
 
 
 @spaces.GPU(duration=16)
-def fill_image(prompt, image, model_selection):
+def fill_image(prompt, image, model_selection, paste_back):
     (
         prompt_embeds,
         negative_prompt_embeds,
@@ -71,8 +71,14 @@ def fill_image(prompt, image, model_selection):
     ):
         yield image, cnet_image
 
-    image = image.convert("RGBA")
-    cnet_image.paste(image, (0, 0), binary_mask)
+    print(f"{model_selection=}")
+    print(f"{paste_back=}")
+
+    if paste_back:
+        image = image.convert("RGBA")
+        cnet_image.paste(image, (0, 0), binary_mask)
+    else:
+        cnet_image = image
 
     yield source, cnet_image
 
@@ -101,14 +107,17 @@ with gr.Blocks() as demo:
                 value="RealVisXL V5.0 Lightning",
                 label="Model",
             )
-            run_button = gr.Button("Generate")
+
+            with gr.Row():
+                with gr.Column():
+                    run_button = gr.Button("Generate")
+
+                with gr.Column():
+                    paste_back = gr.Checkbox(True, label="Paste back original")
 
     with gr.Row():
         input_image = gr.ImageMask(
-            type="pil",
-            label="Input Image",
-            crop_size=(1024, 1024),
-            layers=False
+            type="pil", label="Input Image", crop_size=(1024, 1024), layers=False
         )
 
         result = ImageSlider(
@@ -116,15 +125,13 @@ with gr.Blocks() as demo:
             label="Generated Image",
         )
 
-    use_as_input_button = gr.Button("Use as Input Image", visible=False)  
+    use_as_input_button = gr.Button("Use as Input Image", visible=False)
 
     def use_output_as_input(output_image):
         return gr.update(value=output_image[1])
 
     use_as_input_button.click(
-        fn=use_output_as_input,
-        inputs=[result],
-        outputs=[input_image]
+        fn=use_output_as_input, inputs=[result], outputs=[input_image]
     )
 
     run_button.click(
@@ -137,7 +144,7 @@ with gr.Blocks() as demo:
         outputs=use_as_input_button,
     ).then(
         fn=fill_image,
-        inputs=[prompt, input_image, model_selection],
+        inputs=[prompt, input_image, model_selection, paste_back],
         outputs=result,
     ).then(
         fn=lambda: gr.update(visible=True),
@@ -155,7 +162,7 @@ with gr.Blocks() as demo:
         outputs=use_as_input_button,
     ).then(
         fn=fill_image,
-        inputs=[prompt, input_image, model_selection],
+        inputs=[prompt, input_image, model_selection, paste_back],
         outputs=result,
     ).then(
         fn=lambda: gr.update(visible=True),
@@ -164,4 +171,4 @@ with gr.Blocks() as demo:
     )
 
 
-demo.launch(share=False)
+demo.launch(share=False)