Update app.py

app.py

@@ -33,15 +33,22 @@ 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",
-).to("cuda")
-
-pipe.scheduler = TCDScheduler.from_config(pipe.scheduler.config)
+# --- Define available pipelines ---
+model_mapping = {
+    "RealVisXL V5.0 Lightning": "SG161222/RealVisXL_V5.0_Lightning",
+    "RealVisXL V4.0 Lightning": "SG161222/RealVisXL_V4.0_Lightning",
+}
+pipelines = {}
+for name, repo in model_mapping.items():
+    pipe = StableDiffusionXLFillPipeline.from_pretrained(
+        repo,
+        torch_dtype=torch.float16,
+        vae=vae,
+        controlnet=model,
+        variant="fp16",
+    ).to("cuda")
+    pipe.scheduler = TCDScheduler.from_config(pipe.scheduler.config)
+    pipelines[name] = pipe
 
 def prepare_image_and_mask(image, width, height, overlap_percentage, resize_option, custom_resize_percentage, alignment, overlap_left, overlap_right, overlap_top, overlap_bottom):
     target_size = (width, height)
@@ -237,11 +244,18 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
                     label="Input Image"
                 )
 
+                with gr.Row():
+                    model_selector = gr.Dropdown(
+                    label="Select Model",
+                    choices=list(pipelines.keys()),
+                    value="RealVisXL V5.0 Lightning",
+                )
+                
                 with gr.Row():
                     with gr.Column(scale=2):
                         prompt_input = gr.Textbox(label="Prompt (Optional)")
                     with gr.Column(scale=1):
-                        run_button = gr.Button("Generate")
+                        run_button = gr.Button("Outpaint Image")
 
                 with gr.Row():
                     target_ratio = gr.Radio(
@@ -351,15 +365,17 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
         queue=False
     )
     
-    run_button.click(
+     run_button.click(
         fn=clear_result,
         inputs=None,
         outputs=result,
     ).then(
         fn=infer,
-        inputs=[input_image, 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],
+        inputs=[input_image, 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,
+                model_selector],
         outputs=result,
     ).then(
         fn=lambda x, history: update_history(x, history),
@@ -373,9 +389,11 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
         outputs=result,
     ).then(
         fn=infer,
-        inputs=[input_image, 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],
+        inputs=[input_image, 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,
+                model_selector],
         outputs=result,
     ).then(
         fn=lambda x, history: update_history(x, history),