Update app.py

app.py

@@ -1,3 +1,4 @@
+import io
 import gradio as gr
 import spaces
 import torch
@@ -12,6 +13,15 @@ from pipeline_fill_sd_xl import StableDiffusionXLFillPipeline
 from PIL import Image, ImageDraw
 import numpy as np
 
+# --- NEW: FastAPI bits for custom REST endpoint ---
+from fastapi import FastAPI, File, UploadFile, Form
+from fastapi.responses import StreamingResponse, JSONResponse
+from fastapi.middleware.cors import CORSMiddleware
+# -------------------------------------------------
+
+# =========================
+#   MODEL / PIPELINE LOAD
+# =========================
 config_file = hf_hub_download(
     "xinsir/controlnet-union-sdxl-1.0",
     filename="config_promax.json",
@@ -20,26 +30,19 @@ config_file = hf_hub_download(
 config = ControlNetModel_Union.load_config(config_file)
 controlnet_model = ControlNetModel_Union.from_config(config)
 
-# Load the state dictionary
 model_file = hf_hub_download(
     "xinsir/controlnet-union-sdxl-1.0",
     filename="diffusion_pytorch_model_promax.safetensors",
 )
 state_dict = load_state_dict(model_file)
-
-# Extract the keys from the state_dict
 loaded_keys = list(state_dict.keys())
 
-# Call the method and store all returns in a variable
 result = ControlNetModel_Union._load_pretrained_model(
     controlnet_model, state_dict, model_file, "xinsir/controlnet-union-sdxl-1.0", loaded_keys
 )
-
-# Use the first element from the result
 model = result[0]
 model = model.to(device="cuda", dtype=torch.float16)
 
-
 vae = AutoencoderKL.from_pretrained(
     "madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16
 ).to("cuda")
@@ -55,8 +58,10 @@ pipe = StableDiffusionXLFillPipeline.from_pretrained(
 pipe.scheduler = TCDScheduler.from_config(pipe.scheduler.config)
 
 
+# =========================
+#         HELPERS
+# =========================
 def can_expand(source_width, source_height, target_width, target_height, alignment):
-    """Checks if the image can be expanded based on the alignment."""
     if alignment in ("Left", "Right") and source_width >= target_width:
         return False
     if alignment in ("Top", "Bottom") and source_height >= target_height:
@@ -66,15 +71,13 @@ def can_expand(source_width, source_height, target_width, target_height, alignme
 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)
 
-    # Calculate the scaling factor to fit the image within the target size
+    # Fit image into target canvas
     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)
-    
-    # Resize the source image to fit within target size
     source = image.resize((new_width, new_height), Image.LANCZOS)
 
-    # Apply resize option using percentages
+    # Resize option (%)
     if resize_option == "Full":
         resize_percentage = 100
     elif resize_option == "50%":
@@ -83,66 +86,44 @@ def prepare_image_and_mask(image, width, height, overlap_percentage, resize_opti
         resize_percentage = 33
     elif resize_option == "25%":
         resize_percentage = 25
-    else:  # Custom
+    else:
         resize_percentage = custom_resize_percentage
 
-    # Calculate new dimensions based on percentage
-    resize_factor = resize_percentage / 100
-    new_width = int(source.width * resize_factor)
-    new_height = int(source.height * resize_factor)
-
-    # Ensure minimum size of 64 pixels
-    new_width = max(new_width, 64)
-    new_height = max(new_height, 64)
-
-    # Resize the image
+    resize_factor = max(1, int(resize_percentage)) / 100.0
+    new_width = max(int(source.width * resize_factor), 64)
+    new_height = max(int(source.height * resize_factor), 64)
     source = source.resize((new_width, new_height), Image.LANCZOS)
 
-    # Calculate the overlap in pixels based on the percentage
-    overlap_x = int(new_width * (overlap_percentage / 100))
-    overlap_y = int(new_height * (overlap_percentage / 100))
+    overlap_x = max(int(new_width * (overlap_percentage / 100)), 1)
+    overlap_y = max(int(new_height * (overlap_percentage / 100)), 1)
 
-    # Ensure minimum overlap of 1 pixel
-    overlap_x = max(overlap_x, 1)
-    overlap_y = max(overlap_y, 1)
-
-    # Calculate margins based on alignment
     if alignment == "Middle":
         margin_x = (target_size[0] - new_width) // 2
         margin_y = (target_size[1] - new_height) // 2
     elif alignment == "Left":
-        margin_x = 0
-        margin_y = (target_size[1] - new_height) // 2
+        margin_x = 0; margin_y = (target_size[1] - new_height) // 2
     elif alignment == "Right":
-        margin_x = target_size[0] - new_width
-        margin_y = (target_size[1] - new_height) // 2
+        margin_x = target_size[0] - new_width; margin_y = (target_size[1] - new_height) // 2
     elif alignment == "Top":
-        margin_x = (target_size[0] - new_width) // 2
-        margin_y = 0
+        margin_x = (target_size[0] - new_width) // 2; margin_y = 0
     elif alignment == "Bottom":
-        margin_x = (target_size[0] - new_width) // 2
-        margin_y = target_size[1] - new_height
+        margin_x = (target_size[0] - new_width) // 2; margin_y = target_size[1] - new_height
 
-    # Adjust margins to eliminate gaps
     margin_x = max(0, min(margin_x, target_size[0] - new_width))
     margin_y = max(0, min(margin_y, target_size[1] - new_height))
 
-    # Create a new background image and paste the resized source image
     background = Image.new('RGB', target_size, (255, 255, 255))
     background.paste(source, (margin_x, margin_y))
 
-    # Create the mask
     mask = Image.new('L', target_size, 255)
     mask_draw = ImageDraw.Draw(mask)
 
-    # Calculate overlap areas
     white_gaps_patch = 2
-
     left_overlap = margin_x + overlap_x if overlap_left else margin_x + white_gaps_patch
     right_overlap = margin_x + new_width - overlap_x if overlap_right else margin_x + new_width - white_gaps_patch
     top_overlap = margin_y + overlap_y if overlap_top else margin_y + white_gaps_patch
     bottom_overlap = margin_y + new_height - overlap_y if overlap_bottom else margin_y + new_height - white_gaps_patch
-    
+
     if alignment == "Left":
         left_overlap = margin_x + overlap_x if overlap_left else margin_x
     elif alignment == "Right":
@@ -152,37 +133,19 @@ def prepare_image_and_mask(image, width, height, overlap_percentage, resize_opti
     elif alignment == "Bottom":
         bottom_overlap = margin_y + new_height - overlap_y if overlap_bottom else margin_y + new_height
 
-
-    # Draw the mask
-    mask_draw.rectangle([
-        (left_overlap, top_overlap),
-        (right_overlap, bottom_overlap)
-    ], fill=0)
-
+    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)
-    
-    # Create a preview image showing the mask
-    preview = background.copy().convert('RGBA')
-    
-    # Create a semi-transparent red overlay
-    red_overlay = Image.new('RGBA', background.size, (255, 0, 0, 64))  # Reduced alpha to 64 (25% opacity)
-    
-    # Convert black pixels in the mask to semi-transparent red
-    red_mask = Image.new('RGBA', background.size, (0, 0, 0, 0))
-    red_mask.paste(red_overlay, (0, 0), mask)
-    
-    # Overlay the red mask on the background
-    preview = Image.alpha_composite(preview, red_mask)
-    
-    return preview
-
+# --- NEW: single-call synchronous generator for both UI and REST ---
 @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 run_outpaint_sync(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"
 
@@ -191,7 +154,6 @@ def infer(image, width, height, overlap_percentage, num_inference_steps, resize_
 
     final_prompt = f"{prompt_input} , high quality, 4k" if prompt_input else "high quality, 4k"
 
-    # Use with torch.autocast to ensure consistent dtype
     with torch.autocast(device_type="cuda", dtype=torch.float16):
         (
             prompt_embeds,
@@ -200,7 +162,8 @@ def infer(image, width, height, overlap_percentage, num_inference_steps, resize_
             negative_pooled_prompt_embeds,
         ) = pipe.encode_prompt(final_prompt, "cuda", True)
 
-        for image in pipe(
+        last_image = None
+        for img in pipe(
             prompt_embeds=prompt_embeds,
             negative_prompt_embeds=negative_prompt_embeds,
             pooled_prompt_embeds=pooled_prompt_embeds,
@@ -208,31 +171,54 @@ def infer(image, width, height, overlap_percentage, num_inference_steps, resize_
             image=cnet_image,
             num_inference_steps=num_inference_steps
         ):
-            yield cnet_image, image
+            last_image = img
+
+    if last_image is None:
+        raise RuntimeError("Pipeline did not return an image.")
 
-    image = image.convert("RGBA")
-    cnet_image.paste(image, (0, 0), mask)
+    last_image = last_image.convert("RGBA")
+    cnet_image.paste(last_image, (0, 0), mask)
+    return background, cnet_image
 
+# (Original streaming infer for UI remains, unchanged)
+@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)
+    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" if prompt_input else "high quality, 4k"
+    with torch.autocast(device_type="cuda", dtype=torch.float16):
+        (
+            prompt_embeds,
+            negative_prompt_embeds,
+            pooled_prompt_embeds,
+            negative_pooled_prompt_embeds,
+        ) = pipe.encode_prompt(final_prompt, "cuda", True)
+        for img 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
+        ):
+            yield cnet_image, img
+    img = img.convert("RGBA")
+    cnet_image.paste(img, (0, 0), mask)
     yield background, cnet_image
 
 def clear_result():
-    """Clears the result ImageSlider."""
     return gr.update(value=None)
 
 def preload_presets(target_ratio, ui_width, ui_height):
-    """Updates the width and height sliders based on the selected aspect ratio."""
     if target_ratio == "9:16":
-        changed_width = 720
-        changed_height = 1280
-        return changed_width, changed_height, gr.update()
+        return 720, 1280, gr.update()
     elif target_ratio == "16:9":
-        changed_width = 1280
-        changed_height = 720
-        return changed_width, changed_height, gr.update()
+        return 1280, 720, gr.update()
     elif target_ratio == "1:1":
-        changed_width = 1024
-        changed_height = 1024
-        return changed_width, changed_height, gr.update()
+        return 1024, 1024, gr.update()
     elif target_ratio == "Custom":
         return ui_width, ui_height, gr.update(open=True)
 
@@ -250,7 +236,6 @@ def toggle_custom_resize_slider(resize_option):
     return gr.update(visible=(resize_option == "Custom"))
 
 def update_history(new_image, history):
-    """Updates the history gallery with the new image."""
     if history is None:
         history = []
     history.insert(0, new_image)
@@ -262,68 +247,37 @@ css = """
 }
 """
 
-# Define the title HTML string
-title = """<h1 align="center">Re-Size Image Outpaint</h1>
-"""
+title = """<h1 align="center">Re-Size Image Outpaint</h1>"""
 
 with gr.Blocks(theme="soft", css=css) as demo:
     with gr.Column():
         gr.HTML(title)
-
         with gr.Row():
             with gr.Column():
-                input_image = gr.Image(
-                    type="pil",
-                    label="Input Image"
-                )
-
+                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
+                        value="9:16", scale=2
                     )
-                    
                     alignment_dropdown = gr.Dropdown(
                         choices=["Middle", "Left", "Right", "Top", "Bottom"],
-                        value="Middle",
-                        label="Alignment"
+                        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)
@@ -331,24 +285,11 @@ with gr.Blocks(theme="soft", css=css) as demo:
                                 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")
-                            
-                            
+
                 gr.Examples(
                     examples=[
                         ["./examples/example_2.jpg", 1440, 810, "Left"],
@@ -358,108 +299,106 @@ with gr.Blocks(theme="soft", css=css) as demo:
                     inputs=[input_image, width_slider, height_slider, alignment_dropdown],
                 )
 
-                
-
             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):
-        """Sets the generated output as the new input 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(
-        # --- FIX APPLIED HERE ---
-        # Safely update history only if the result (x) is not None.
-        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(
-        # --- FIX APPLIED HERE ---
-        # Safely update history only if the result (x) is not None.
-        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,
-    )
+    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,
+        fn=lambda *args: preview_image_and_mask(*args),
         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
+        outputs=preview_image, queue=False
     )
 
-demo.queue(max_size=12).launch(share=False)
+# =========================================
+#     FASTAPI APP + CUSTOM REST ENDPOINT
+# =========================================
+app = FastAPI()
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"], allow_credentials=True,
+    allow_methods=["*"], allow_headers=["*"],
+)
+
+@app.post("/rest/infer")
+def rest_infer(
+    file: UploadFile = File(...),
+    width: int = Form(1024),
+    height: int = Form(1024),
+    overlap_percentage: float = Form(10),
+    num_inference_steps: int = Form(8),
+    resize_option: str = Form("Full"),
+    custom_resize_percentage: float = Form(50),
+    prompt_input: str = Form(""),
+    alignment: str = Form("Middle"),
+    overlap_left: bool = Form(True),
+    overlap_right: bool = Form(True),
+    overlap_top: bool = Form(True),
+    overlap_bottom: bool = Form(True),
+):
+    try:
+        img_bytes = file.file.read()
+        image = Image.open(io.BytesIO(img_bytes)).convert("RGB")
+    except Exception as e:
+        return JSONResponse({"error": f"Invalid image upload: {e}"}, status_code=400)
+
+    try:
+        _, outpainted = run_outpaint_sync(
+            image=image,
+            width=width,
+            height=height,
+            overlap_percentage=overlap_percentage,
+            num_inference_steps=num_inference_steps,
+            resize_option=resize_option,
+            custom_resize_percentage=custom_resize_percentage,
+            prompt_input=prompt_input,
+            alignment=alignment,
+            overlap_left=overlap_left,
+            overlap_right=overlap_right,
+            overlap_top=overlap_top,
+            overlap_bottom=overlap_bottom,
+        )
+    except Exception as e:
+        return JSONResponse({"error": str(e)}, status_code=500)
+
+    buf = io.BytesIO()
+    outpainted.save(buf, format="PNG")
+    buf.seek(0)
+    return StreamingResponse(buf, media_type="image/png")
+
+# Mount the Gradio UI at root path
+app = gr.mount_gradio_app(app, demo, path="/")