Update app.py

app.py

@@ -10,7 +10,6 @@ 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",
@@ -35,68 +34,58 @@ model.to(device="cuda", dtype=torch.float16)
 vae = AutoencoderKL.from_pretrained(
     "madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16
 ).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
+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")
 
 @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)
-    
-    # Encode the prompt and ensure the embeddings are in float16
+
     (
         prompt_embeds,
         negative_prompt_embeds,
         pooled_prompt_embeds,
         negative_pooled_prompt_embeds,
     ) = pipe.encode_prompt(prompt, "cuda", True)
-    prompt_embeds = prompt_embeds.to(torch.float16)
-    negative_prompt_embeds = negative_prompt_embeds.to(torch.float16)
-    pooled_prompt_embeds = pooled_prompt_embeds.to(torch.float16)
-    negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.to(torch.float16)
-    
     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 output_image in pipe(
+
+    for 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 output_image, cnet_image
-    
+        yield image, cnet_image
+
     print(f"{model_selection=}")
     print(f"{paste_back=}")
-
     if paste_back:
-        output_image = output_image.convert("RGBA")
-        cnet_image.paste(output_image, (0, 0), binary_mask)
+        image = image.convert("RGBA")
+        cnet_image.paste(image, (0, 0), binary_mask)
     else:
-        cnet_image = output_image
-    
+        cnet_image = image
     yield source, cnet_image
 
 def clear_result():
@@ -242,47 +231,32 @@ 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, model_selection):
-    if image is None:
-        yield None, None
-        return
-    
+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):
     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"
-    
-    pipe = get_pipeline(model_selection)
-    # Encode the prompt and ensure the embeddings are in float16
+    final_prompt = f"{prompt_input} , high quality, 4k"
     (
         prompt_embeds,
         negative_prompt_embeds,
         pooled_prompt_embeds,
         negative_pooled_prompt_embeds,
     ) = pipe.encode_prompt(final_prompt, "cuda", True)
-    prompt_embeds = prompt_embeds.to(torch.float16)
-    negative_prompt_embeds = negative_prompt_embeds.to(torch.float16)
-    pooled_prompt_embeds = pooled_prompt_embeds.to(torch.float16)
-    negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.to(torch.float16)
-    
-    for output_image in pipe(
+    for 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, output_image
-    
-    output_image = output_image.convert("RGBA")
-    cnet_image.paste(output_image, (0, 0), mask)
+        yield cnet_image, image
+    image = image.convert("RGBA")
+    cnet_image.paste(image, (0, 0), mask)
     yield background, cnet_image
 
-
 def use_output_as_input(output_image):
     return gr.update(value=output_image[1])
 
@@ -492,11 +466,6 @@ 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=[
@@ -515,30 +484,35 @@ 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],
@@ -552,7 +526,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, out_model_selection],  # Add model_selection here
+                overlap_left, overlap_right, overlap_top, overlap_bottom],
         outputs=[result_outpaint],
     ).then(
         fn=lambda x, history: update_history(x[1], history),
@@ -571,7 +545,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, out_model_selection],  # Add model_selection here
+                overlap_left, overlap_right, overlap_top, overlap_bottom],
         outputs=[result_outpaint],
     ).then(
         fn=lambda x, history: update_history(x[1], history),
@@ -589,4 +563,5 @@ with gr.Blocks(css=css, fill_height=True) as demo:
         outputs=[preview_image],
         queue=False
     )
+
 demo.launch(show_error=True)