Update app.py

app.py

@@ -18,6 +18,8 @@ config_file = hf_hub_download(
     filename="config_promax.json",
 )
 
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
 config = ControlNetModel_Union.load_config(config_file)
 controlnet_model = ControlNetModel_Union.from_config(config)
 model_file = hf_hub_download(
@@ -28,11 +30,10 @@ state_dict = load_state_dict(model_file)
 model, _, _, _, _ = ControlNetModel_Union._load_pretrained_model(
     controlnet_model, state_dict, model_file, "xinsir/controlnet-union-sdxl-1.0"
 )
-device = 'cuda' if torch.cuda.is_available() else 'cpu'
-
 model.to(device=device, dtype=torch.float16 if device == 'cuda' else torch.float32)
+
 vae = AutoencoderKL.from_pretrained(
-    "madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16 if device == 'cuda' else torch.float32
+    "madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16 if device == 'cuda' else torch.float32,
 ).to(device)
 
 pipe = StableDiffusionXLFillPipeline.from_pretrained(
@@ -40,21 +41,22 @@ pipe = StableDiffusionXLFillPipeline.from_pretrained(
     torch_dtype=torch.float16 if device == 'cuda' else torch.float32,
     vae=vae,
     controlnet=model,
-    variant="fp16" if device == 'cuda' else None,
+    #variant="fp16",
 ).to(device)
 
 pipe.scheduler = TCDScheduler.from_config(pipe.scheduler.config)
+if device == 'cuda':
+    pipe.enable_model_cpu_offload()
+    pipe.enable_xformers_memory_efficient_attention()
+
+def fill_image(prompt, image, model_selection, paste_back):
+    (
+        prompt_embeds,
+        negative_prompt_embeds,
+        pooled_prompt_embeds,
+        negative_pooled_prompt_embeds,
+    ) = pipe.encode_prompt(prompt, "cuda", True)
 
-prompt = "high quality"
-(
-    prompt_embeds,
-    negative_prompt_embeds,
-    pooled_prompt_embeds,
-    negative_pooled_prompt_embeds,
-) = pipe.encode_prompt(prompt, device, True)
-
-
-def fill_image(image, model_selection):
     source = image["background"]
     mask = image["layers"][0]
 
@@ -72,8 +74,14 @@ def fill_image(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
 
@@ -82,30 +90,33 @@ def clear_result():
     return gr.update(value=None)
 
 
-css = """
-.gradio-container {
-    width: 1024px !important;
-}
-"""
-
-
-title = """<h1 align="center">Diffusers Image Fill</h1>
-<div align="center">Draw the mask over the subject you want to erase or change.</div>
-<div align="center">This space is a PoC made for the guide <a href='https://huggingface.co/blog/OzzyGT/diffusers-image-fill'>Diffusers Image Fill</a>.</div>
-"""
-
-with gr.Blocks(css=css) as demo:
-    gr.HTML(title)
-
-    run_button = gr.Button("Generate")
+with gr.Blocks() as demo:
+    gr.Markdown('# Ai object removal')
+    with gr.Row():
+        with gr.Column():
+            prompt = gr.Textbox(
+                label="Prompt",
+                info="Describe what to inpaint the mask with",
+                lines=3,
+            )
+        with gr.Column():
+            model_selection = gr.Dropdown(
+                choices=list(MODELS.keys()),
+                value="RealVisXL V5.0 Lightning",
+                label="Model",
+                visible = False,
+            )
+
+            with gr.Row():
+                with gr.Column():
+                    run_button = gr.Button("Generate")
+
+                with gr.Column():
+                    paste_back = gr.Checkbox(False, label="Paste back original", visible = False)
 
     with gr.Row():
         input_image = gr.ImageMask(
-            type="pil",
-            label="Input Image",
-            canvas_size=(1024, 1024),
-            layers=False,
-            sources=["upload"],
+            type="pil", label="Input Image", layers=False
         )
 
         result = ImageSlider(
@@ -113,21 +124,50 @@ with gr.Blocks(css=css) as demo:
             label="Generated Image",
         )
 
-    model_selection = gr.Dropdown(
-        choices=list(MODELS.keys()),
-        value="RealVisXL V5.0 Lightning",
-        label="Model",
+    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]
     )
 
     run_button.click(
         fn=clear_result,
         inputs=None,
         outputs=result,
+    ).then(
+        fn=lambda: gr.update(visible=False),
+        inputs=None,
+        outputs=use_as_input_button,
+    ).then(
+        fn=fill_image,
+        inputs=[prompt, input_image, model_selection, paste_back],
+        outputs=result,
+    ).then(
+        fn=lambda: gr.update(visible=True),
+        inputs=None,
+        outputs=use_as_input_button,
+    )
+
+    prompt.submit(
+        fn=clear_result,
+        inputs=None,
+        outputs=result,
+    ).then(
+        fn=lambda: gr.update(visible=False),
+        inputs=None,
+        outputs=use_as_input_button,
     ).then(
         fn=fill_image,
-        inputs=[input_image, model_selection],
+        inputs=[prompt, input_image, model_selection, paste_back],
         outputs=result,
+    ).then(
+        fn=lambda: gr.update(visible=True),
+        inputs=None,
+        outputs=use_as_input_button,
     )
 
 
-demo.launch(share=False)
+demo.launch(share=True,debug = True)