wip outpainting

app.py

@@ -4,6 +4,9 @@ import torch
 from loadimg import load_img
 from torchvision import transforms
 from transformers import AutoModelForImageSegmentation
+from diffusers import FluxFillPipeline
+from PIL import Image, ImageDraw
+from diffusers.utils import load_image
 
 torch.set_float32_matmul_precision(["high", "highest"][0])
 
@@ -20,10 +23,186 @@ transform_image = transforms.Compose(
     ]
 )
 
+pipe = FluxFillPipeline.from_pretrained(
+    "black-forest-labs/FLUX.1-Fill-dev", torch_dtype=torch.bfloat16
+).to("cuda")
+
+
+def can_expand(source_width, source_height, target_width, target_height, alignment):
+    if alignment in ("Left", "Right") and source_width >= target_width:
+        return False
+    if alignment in ("Top", "Bottom") and source_height >= target_height:
+        return False
+    return True
+
+
+def prepare_image_and_mask(
+    image,
+    width,
+    height,
+    overlap_percentage,
+    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)
+
+    source = image.resize((new_width, new_height), Image.LANCZOS)
+
+    resize_percentage = 50
+
+    # 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
+    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))
+
+    # 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
+    elif alignment == "Right":
+        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
+    elif alignment == "Bottom":
+        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":
+        right_overlap = (
+            margin_x + new_width - overlap_x if overlap_right else margin_x + new_width
+        )
+    elif alignment == "Top":
+        top_overlap = margin_y + overlap_y if overlap_top else margin_y
+    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
+    )
+
+    return background, mask
+
+
+def inpaint(
+    image,
+    width,
+    height,
+    overlap_percentage,
+    num_inference_steps,
+    custom_resize_percentage,
+    prompt_input,
+    alignment,
+    overlap_left,
+    overlap_right,
+    overlap_top,
+    overlap_bottom,
+    progress=gr.Progress(track_tqdm=True),
+):
+    background, mask = prepare_image_and_mask(
+        image,
+        width,
+        height,
+        overlap_percentage,
+        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 = prompt_input
+
+    # generator = torch.Generator(device="cuda").manual_seed(42)
+
+    result = pipe(
+        prompt=final_prompt,
+        height=height,
+        width=width,
+        image=cnet_image,
+        mask_image=mask,
+        num_inference_steps=num_inference_steps,
+        guidance_scale=30,
+    ).images[0]
+
+    result = result.convert("RGBA")
+    cnet_image.paste(result, (0, 0), mask)
+
+    return cnet_image
+
 
 @spaces.GPU
-def rmbg(image,url):
-    if image is None : 
+def rmbg(image, url):
+    if image is None:
         image = url
     image = load_img(image).convert("RGB")
     image_size = image.size
@@ -38,11 +217,21 @@ def rmbg(image,url):
     return image
 
 
-rmbg_tab = gr.Interface(fn=rmbg, inputs=["image","text"], outputs=["image"], api_name="rmbg")
+def placeholder(img):
+    return img
+
+
+rmbg_tab = gr.Interface(
+    fn=rmbg, inputs=["image", "text"], outputs=["image"], api_name="rmbg"
+)
+
+outpaint_tab = gr.Interface(
+    fr=placeholder, inputs=["image"], outputs=["image"], api_name="outpainting"
+)
 
 demo = gr.TabbedInterface(
-    [rmbg_tab],
-    ["remove background"],
+    [rmbg_tab, outpaint_tab],
+    ["remove background", "outpainting"],
     title="Utilities that require GPU",
 )
 

requirements.txt

@@ -11,3 +11,4 @@ scikit-image
 kornia
 transformers
 huggingface_hub
+diffusers