Update app.py

app.py

@@ -11,31 +11,20 @@ from pipeline_fill_sd_xl import StableDiffusionXLFillPipeline
 
 from PIL import Image, ImageDraw
 
-# =========================
-#   LOAD MODELS (original)
-# =========================
-config_file = hf_hub_download(
-    "xinsir/controlnet-union-sdxl-1.0",
-    filename="config_promax.json",
-)
+# ===== Load models (original from your Space) =====
+config_file = hf_hub_download("xinsir/controlnet-union-sdxl-1.0", filename="config_promax.json")
 config = ControlNetModel_Union.load_config(config_file)
 controlnet_model = ControlNetModel_Union.from_config(config)
 
-model_file = hf_hub_download(
-    "xinsir/controlnet-union-sdxl-1.0",
-    filename="diffusion_pytorch_model_promax.safetensors",
-)
+model_file = hf_hub_download("xinsir/controlnet-union-sdxl-1.0", filename="diffusion_pytorch_model_promax.safetensors")
 state_dict = load_state_dict(model_file)
 loaded_keys = list(state_dict.keys())
-
 result = ControlNetModel_Union._load_pretrained_model(
     controlnet_model, state_dict, model_file, "xinsir/controlnet-union-sdxl-1.0", loaded_keys
 )
 model = result[0].to(device="cuda", dtype=torch.float16)
 
-vae = AutoencoderKL.from_pretrained(
-    "madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16
-).to("cuda")
+vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16).to("cuda")
 
 pipe = StableDiffusionXLFillPipeline.from_pretrained(
     "SG161222/RealVisXL_V5.0_Lightning",
@@ -47,9 +36,7 @@ pipe = StableDiffusionXLFillPipeline.from_pretrained(
 
 pipe.scheduler = TCDScheduler.from_config(pipe.scheduler.config)
 
-# =========================
-#          HELPERS
-# =========================
+# ===== Helpers (original) =====
 def can_expand(source_width, source_height, target_width, target_height, alignment):
     if alignment in ("Left", "Right") and source_width >= target_width:
         return False
@@ -57,9 +44,9 @@ def can_expand(source_width, source_height, target_width, target_height, alignme
         return False
     return True
 
-def prepare_image_and_mask(image, width, height, overlap_percentage, resize_option, custom_resize_percentage, alignment, overlap_left, overlap_right, overlap_top, overlap_bottom):
+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)
-
     scale_factor = min(target_size[0] / image.width, target_size[1] / image.height)
     new_width = int(image.width * scale_factor)
     new_height = int(image.height * scale_factor)
@@ -123,18 +110,24 @@ def prepare_image_and_mask(image, width, height, overlap_percentage, resize_opti
     mask_draw.rectangle([(left_overlap, top_overlap), (right_overlap, bottom_overlap)], fill=0)
     return background, mask
 
-def preview_image_and_mask(image, width, height, overlap_percentage, resize_option, custom_resize_percentage, 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)
+def preview_image_and_mask(image, width, height, overlap_percentage, resize_option, custom_resize_percentage,
+                           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)
     preview = background.copy().convert('RGBA')
     red_overlay = Image.new('RGBA', background.size, (255, 0, 0, 64))
     red_mask = Image.new('RGBA', background.size, (0, 0, 0, 0))
     red_mask.paste(red_overlay, (0, 0), mask)
     return Image.alpha_composite(preview, red_mask)
 
-# ===== Streaming infer for the UI (original) =====
+# ===== Streaming infer (UI) =====
 @spaces.GPU(duration=24)
-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)
+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"
 
@@ -165,14 +158,15 @@ def infer(image, width, height, overlap_percentage, num_inference_steps, resize_
     cnet_image.paste(image, (0, 0), mask)
     yield background, cnet_image
 
-# ===== Non-streaming wrapper used by REST Interface =====
-def infer_rest(image, width, height, overlap_percentage, num_inference_steps, resize_option, custom_resize_percentage, prompt_input, alignment, overlap_left, overlap_right, overlap_top, overlap_bottom):
-    # Run generator to completion and return final pair
-    gen = 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)
+# ===== Non-streaming wrapper (returns final pair) =====
+def infer_rest(image, width, height, overlap_percentage, num_inference_steps, resize_option, custom_resize_percentage,
+               prompt_input, alignment, overlap_left, overlap_right, overlap_top, overlap_bottom):
+    gen = 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)
     last = None
     for last in gen:
         pass
-    return last  # (background, outpainted)
+    return last  # (background, generated)
 
 def clear_result():
     return gr.update(value=None)
@@ -207,69 +201,33 @@ def update_history(new_image, history):
     return history
 
 css = """
-.gradio-container {
-    width: 1200px !important;
-}
+.gradio-container { width: 1200px !important; }
 """
-
 title = """<h1 align="center">Re-Size Image Outpaint</h1>"""
 
-# ---- UI (original Blocks) ----
+# ---- Full UI (unchanged) ----
 with gr.Blocks(theme="soft", css=css) as ui_app:
     with gr.Column():
         gr.HTML(title)
-
         with gr.Row():
             with gr.Column():
                 input_image = gr.Image(type="pil", label="Input Image")
-
                 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")
-
                 with gr.Row():
-                    target_ratio = gr.Radio(
-                        label="Expected Ratio",
-                        choices=["9:16", "16:9", "1:1", "Custom"],
-                        value="9:16",
-                        scale=2
-                    )
-
-                    alignment_dropdown = gr.Dropdown(
-                        choices=["Middle", "Left", "Right", "Top", "Bottom"],
-                        value="Middle",
-                        label="Alignment"
-                    )
-
+                    target_ratio = gr.Radio(label="Expected Ratio", choices=["9:16", "16:9", "1:1", "Custom"], value="9:16", scale=2)
+                    alignment_dropdown = gr.Dropdown(choices=["Middle", "Left", "Right", "Top", "Bottom"], value="Middle", label="Alignment")
                 with gr.Accordion(label="Advanced settings", open=False) as settings_panel:
                     with gr.Column():
                         with gr.Row():
-                            width_slider = gr.Slider(
-                                label="Target Width",
-                                minimum=720,
-                                maximum=1536,
-                                step=8,
-                                value=720,
-                            )
-                            height_slider = gr.Slider(
-                                label="Target Height",
-                                minimum=720,
-                                maximum=1536,
-                                step=8,
-                                value=1280,
-                            )
-
+                            width_slider = gr.Slider(label="Target Width", minimum=720, maximum=1536, step=8, value=720)
+                            height_slider = gr.Slider(label="Target Height", minimum=720, maximum=1536, step=8, value=1280)
                         num_inference_steps = gr.Slider(label="Steps", minimum=4, maximum=12, step=1, value=8)
                         with gr.Group():
-                            overlap_percentage = gr.Slider(
-                                label="Mask overlap (%)",
-                                minimum=1,
-                                maximum=50,
-                                value=10,
-                                step=1
-                            )
+                            overlap_percentage = gr.Slider(label="Mask overlap (%)", minimum=1, maximum=50, value=10, step=1)
                             with gr.Row():
                                 overlap_top = gr.Checkbox(label="Overlap Top", value=True)
                                 overlap_right = gr.Checkbox(label="Overlap Right", value=True)
@@ -277,20 +235,8 @@ with gr.Blocks(theme="soft", css=css) as ui_app:
                                 overlap_left = gr.Checkbox(label="Overlap Left", value=True)
                                 overlap_bottom = gr.Checkbox(label="Overlap Bottom", value=True)
                         with gr.Row():
-                            resize_option = gr.Radio(
-                                label="Resize input image",
-                                choices=["Full", "50%", "33%", "25%", "Custom"],
-                                value="Full"
-                            )
-                            custom_resize_percentage = gr.Slider(
-                                label="Custom resize (%)",
-                                minimum=1,
-                                maximum=100,
-                                step=1,
-                                value=50,
-                                visible=False
-                            )
-
+                            resize_option = gr.Radio(label="Resize input image", choices=["Full", "50%", "33%", "25%", "Custom"], value="Full")
+                            custom_resize_percentage = gr.Slider(label="Custom resize (%)", minimum=1, maximum=100, step=1, value=50, visible=False)
                         with gr.Column():
                             preview_button = gr.Button("Preview alignment and mask")
 
@@ -304,101 +250,46 @@ with gr.Blocks(theme="soft", css=css) as ui_app:
                 )
 
             with gr.Column():
-                result = ImageSlider(
-                    interactive=False,
-                    label="Generated Image",
-                )
+                result = ImageSlider(interactive=False, label="Generated Image")
                 use_as_input_button = 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")
 
     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]
-    )
-
-    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
-    )
-
-    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],
-        outputs=result,
-    ).then(
-        fn=lambda x, history: update_history(x[1], history) if x else history,
-        inputs=[result, history_gallery],
-        outputs=history_gallery,
-    ).then(
-        fn=lambda: gr.update(visible=True),
-        inputs=None,
-        outputs=use_as_input_button,
-    )
-
-    prompt_input.submit(
-        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],
-        outputs=result,
-    ).then(
-        fn=lambda x, history: update_history(x[1], history) if x else history,
-        inputs=[result, history_gallery],
-        outputs=history_gallery,
-    ).then(
-        fn=lambda: gr.update(visible=True),
-        inputs=None,
-        outputs=use_as_input_button,
-    )
-
-    preview_button.click(
-        fn=preview_image_and_mask,
-        inputs=[input_image, width_slider, height_slider, overlap_percentage, resize_option, custom_resize_percentage, alignment_dropdown,
-                overlap_left, overlap_right, overlap_top, overlap_bottom],
-        outputs=preview_image,
-        queue=False
-    )
-
-# ---- API (minimal Interface to guarantee REST route) ----
+    use_as_input_button.click(fn=use_output_as_input, inputs=[result], outputs=[input_image])
+
+    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)
+
+    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],
+              outputs=result) \
+        .then(fn=lambda x, history: update_history(x[1], history) if x else history, inputs=[result, history_gallery], outputs=history_gallery) \
+        .then(fn=lambda: gr.update(visible=True), inputs=None, outputs=use_as_input_button)
+
+    prompt_input.submit(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],
+              outputs=result) \
+        .then(fn=lambda x, history: update_history(x[1], history) if x else history, inputs=[result, history_gallery], outputs=history_gallery) \
+        .then(fn=lambda: gr.update(visible=True), inputs=None, outputs=use_as_input_button)
+
+    preview_button.click(fn=preview_image_and_mask,
+                         inputs=[input_image, width_slider, height_slider, overlap_percentage, resize_option,
+                                 custom_resize_percentage, alignment_dropdown, overlap_left, overlap_right,
+                                 overlap_top, overlap_bottom],
+                         outputs=preview_image, queue=False)
+
+# ---- Minimal Interface tab that DEFINITELY exposes /api/predict/infer ----
 api_app = gr.Interface(
     fn=infer_rest,
     inputs=[
@@ -416,18 +307,15 @@ api_app = gr.Interface(
         gr.Checkbox(value=True, label="Overlap Top"),
         gr.Checkbox(value=True, label="Overlap Bottom"),
     ],
-    outputs=[
-        gr.Image(label="Background"),
-        gr.Image(label="Generated"),
-    ],
+    outputs=[gr.Image(label="Background"), gr.Image(label="Generated")],
     allow_flagging="never",
-    api_name="infer",   # <-- this creates /api/predict/infer
+    api_name="infer",  # <--- THIS creates /api/predict/infer
     title="Re-Size Image Outpaint API",
     description="Non-streaming endpoint for programmatic access.",
 )
 
-# ---- Publish both (UI + API) together ----
-demo = gr.TabbedInterface([ui_app, api_app], tab_names=["App", "API"])
+# Publish BOTH tabs — put API FIRST to be extra safe on older Gradio builds
+demo = gr.TabbedInterface([api_app, ui_app], tab_names=["API", "App"])
 
-# IMPORTANT: expose REST
+# Open REST API
 demo.queue(max_size=12, api_open=True).launch(share=False)