Update app.py

app.py

@@ -10,6 +10,7 @@ from pipeline_fill_sd_xl import StableDiffusionXLFillPipeline
 from PIL import Image, ImageDraw
 import numpy as np
 
+
 MODELS = {
     "RealVisXL V5.0 Lightning": "SG161222/RealVisXL_V5.0_Lightning",
     "Lustify Lightning": "GraydientPlatformAPI/lustify-lightning",
@@ -34,58 +35,63 @@ model.to(device="cuda", dtype=torch.float16)
 vae = AutoencoderKL.from_pretrained(
     "madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16
 ).to("cuda")
-pipe = StableDiffusionXLFillPipeline.from_pretrained(
-    "SG161222/RealVisXL_V5.0_Lightning",
-    torch_dtype=torch.float16,
-    vae=vae,
-    controlnet=model,
-    variant="fp16",
-)
-pipe = StableDiffusionXLFillPipeline.from_pretrained(
-    "GraydientPlatformAPI/lustify-lightning",
-    torch_dtype=torch.float16,
-    vae=vae,
-    controlnet=model,
-)
-pipe.scheduler = TCDScheduler.from_config(pipe.scheduler.config)
-pipe.to("cuda")
+
+pipe = None
+
+
+def get_pipeline(model_name):
+    global pipe
+    if pipe is not None and pipe.config.model_name == MODELS[model_name]:
+        return pipe
+    pipe = StableDiffusionXLFillPipeline.from_pretrained(
+        MODELS[model_name],
+        torch_dtype=torch.float16,
+        vae=vae,
+        controlnet=model,
+    ).to("cuda")
+    pipe.scheduler = TCDScheduler.from_config(pipe.scheduler.config)
+    return pipe
 
 @spaces.GPU(duration=12)
 def fill_image(prompt, image, model_selection, paste_back):
-    print(f"Received image: {image}")
     if image is None:
         yield None, None
         return
-
+    
+    pipe = get_pipeline(model_selection)
+    
     (
         prompt_embeds,
         negative_prompt_embeds,
         pooled_prompt_embeds,
         negative_pooled_prompt_embeds,
     ) = pipe.encode_prompt(prompt, "cuda", True)
+    
     source = image["background"]
     mask = image["layers"][0]
     alpha_channel = mask.split()[3]
     binary_mask = alpha_channel.point(lambda p: p > 0 and 255)
+    
     cnet_image = source.copy()
     cnet_image.paste(0, (0, 0), binary_mask)
-
-    for image in pipe(
+    for output_image in pipe(
         prompt_embeds=prompt_embeds,
         negative_prompt_embeds=negative_prompt_embeds,
         pooled_prompt_embeds=pooled_prompt_embeds,
         negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
         image=cnet_image,
     ):
-        yield image, cnet_image
-
+        yield output_image, cnet_image
+    
     print(f"{model_selection=}")
     print(f"{paste_back=}")
+
     if paste_back:
-        image = image.convert("RGBA")
-        cnet_image.paste(image, (0, 0), binary_mask)
+        output_image = output_image.convert("RGBA")
+        cnet_image.paste(output_image, (0, 0), binary_mask)
     else:
-        cnet_image = image
+        cnet_image = output_image
+    
     yield source, cnet_image
 
 def clear_result():
@@ -231,32 +237,42 @@ def outpaint(image, width, height, overlap_percentage, num_inference_steps, resi
     yield background, cnet_image
 
 @spaces.GPU(duration=12)
-def infer(image, width, height, overlap_percentage, num_inference_steps, resize_option, custom_resize_percentage, prompt_input, alignment, overlap_left, overlap_right, overlap_top, overlap_bottom):
+def infer(image, width, height, overlap_percentage, num_inference_steps, resize_option, custom_resize_percentage, prompt_input, alignment, overlap_left, overlap_right, overlap_top, overlap_bottom, model_selection):
+    if image is None:
+        yield None, None
+        return
+    
     background, mask = prepare_image_and_mask(image, width, height, overlap_percentage, resize_option, custom_resize_percentage, alignment, overlap_left, overlap_right, overlap_top, overlap_bottom)
     if not can_expand(background.width, background.height, width, height, alignment):
         alignment = "Middle"
+    
     cnet_image = background.copy()
     cnet_image.paste(0, (0, 0), mask)
-    final_prompt = f"{prompt_input} , high quality, 4k"
+    final_prompt = f"{prompt_input}, high quality, 4k"
+    
+    pipe = get_pipeline(model_selection)
     (
         prompt_embeds,
         negative_prompt_embeds,
         pooled_prompt_embeds,
         negative_pooled_prompt_embeds,
     ) = pipe.encode_prompt(final_prompt, "cuda", True)
-    for image in pipe(
+    
+    for output_image in pipe(
         prompt_embeds=prompt_embeds,
         negative_prompt_embeds=negative_prompt_embeds,
         pooled_prompt_embeds=pooled_prompt_embeds,
         negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
         image=cnet_image,
-        num_inference_steps=num_inference_steps
+        num_inference_steps=num_inference_steps,
     ):
-        yield cnet_image, image
-    image = image.convert("RGBA")
-    cnet_image.paste(image, (0, 0), mask)
+        yield cnet_image, output_image
+    
+    output_image = output_image.convert("RGBA")
+    cnet_image.paste(output_image, (0, 0), mask)
     yield background, cnet_image
 
+
 def use_output_as_input(output_image):
     return gr.update(value=output_image[1])
 
@@ -466,6 +482,11 @@ with gr.Blocks(css=css, fill_height=True) as demo:
                                         visible=False
                                     )
                                 with gr.Column():
+                                    out_model_selection = gr.Dropdown(
+                                        choices=list(MODELS.keys()),
+                                        value="RealVisXL V5.0 Lightning",
+                                        label="Model",
+                                    )
                                     preview_button = gr.Button("Preview alignment and mask")
                         gr.Examples(
                             examples=[
@@ -484,35 +505,30 @@ with gr.Blocks(css=css, fill_height=True) as demo:
                         use_as_input_button_outpaint = gr.Button("Use as Input Image", visible=False)
                         history_gallery = gr.Gallery(label="History", columns=6, object_fit="contain", interactive=False)
                         preview_image = gr.Image(label="Preview")
-
     target_ratio.change(
         fn=preload_presets,
         inputs=[target_ratio, width_slider, height_slider],
         outputs=[width_slider, height_slider, settings_panel],
         queue=False
     )
-
     width_slider.change(
         fn=select_the_right_preset,
         inputs=[width_slider, height_slider],
         outputs=[target_ratio],
         queue=False
     )
-
     height_slider.change(
         fn=select_the_right_preset,
         inputs=[width_slider, height_slider],
         outputs=[target_ratio],
         queue=False
     )
-
     resize_option.change(
         fn=toggle_custom_resize_slider,
         inputs=[resize_option],
         outputs=[custom_resize_percentage],
         queue=False
     )
-                        
     use_as_input_button_outpaint.click(
         fn=use_output_as_input,
         inputs=[result_outpaint],
@@ -526,7 +542,7 @@ with gr.Blocks(css=css, fill_height=True) as demo:
         fn=infer,
         inputs=[input_image_outpaint, width_slider, height_slider, overlap_percentage, num_inference_steps,
                 resize_option, custom_resize_percentage, prompt_input, alignment_dropdown,
-                overlap_left, overlap_right, overlap_top, overlap_bottom],
+                overlap_left, overlap_right, overlap_top, overlap_bottom, out_model_selection],  # Add model_selection here
         outputs=[result_outpaint],
     ).then(
         fn=lambda x, history: update_history(x[1], history),
@@ -545,7 +561,7 @@ with gr.Blocks(css=css, fill_height=True) as demo:
         fn=infer,
         inputs=[input_image_outpaint, width_slider, height_slider, overlap_percentage, num_inference_steps,
                 resize_option, custom_resize_percentage, prompt_input, alignment_dropdown,
-                overlap_left, overlap_right, overlap_top, overlap_bottom],
+                overlap_left, overlap_right, overlap_top, overlap_bottom, out_model_selection],  # Add model_selection here
         outputs=[result_outpaint],
     ).then(
         fn=lambda x, history: update_history(x[1], history),
@@ -563,5 +579,4 @@ with gr.Blocks(css=css, fill_height=True) as demo:
         outputs=[preview_image],
         queue=False
     )
-
 demo.launch(show_error=True)