Spaces:
Paused
Paused
| import os | |
| import torch | |
| import spaces | |
| import gradio as gr | |
| from diffusers import FluxFillPipeline | |
| import random | |
| import numpy as np | |
| from huggingface_hub import hf_hub_download | |
| from PIL import Image, ImageOps | |
| CSS = """ | |
| h1 { | |
| margin-top: 10px | |
| } | |
| """ | |
| os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1" | |
| MAX_SEED = np.iinfo(np.int32).max | |
| repo_id = "black-forest-labs/FLUX.1-Fill-dev" | |
| if torch.cuda.is_available(): | |
| pipe = FluxFillPipeline.from_pretrained(repo_id, torch_dtype=torch.bfloat16).to("cuda") | |
| def gen( | |
| prompt, | |
| image, | |
| mask_image, | |
| width, | |
| height, | |
| num_inference_steps, | |
| seed, | |
| guidance_scale, | |
| ): | |
| generator = torch.Generator("cpu").manual_seed(seed) | |
| result = pipe( | |
| prompt=prompt, | |
| image=image, | |
| mask_image=mask_image, | |
| width=width, | |
| height=height, | |
| num_inference_steps=num_inference_steps, | |
| generator=generator, | |
| guidance_scale=guidance_scale, | |
| max_sequence_length=512, | |
| ).images[0] | |
| return result | |
| def inpaintGen( | |
| imgMask, | |
| inpaint_prompt: str, | |
| guidance: float, | |
| num_steps: int, | |
| seed: int, | |
| randomize_seed: bool, | |
| progress=gr.Progress(track_tqdm=True)): | |
| source_path = imgMask["background"] | |
| mask_path = imgMask["layers"][0] | |
| if not source_path: | |
| raise gr.Error("Please upload an image.") | |
| if not mask_path: | |
| raise gr.Error("Please draw a mask on the image.") | |
| source_img = Image.open(source_path).convert("RGB") | |
| mask_img = Image.open(mask_path) | |
| alpha_channel=mask_img.split()[3] | |
| binary_mask = alpha_channel.point(lambda p: p > 0 and 255) | |
| width, height = source_img.size | |
| new_width = (width // 16) * 16 | |
| new_height = (height // 16) * 16 | |
| # If the image size is not already divisible by 16, resize it | |
| if width != new_width or height != new_height: | |
| source_img = source_img.resize((new_width, new_height), Image.LANCZOS) | |
| if randomize_seed: | |
| seed = random.randint(0, MAX_SEED) | |
| generator = torch.Generator("cpu").manual_seed(seed) | |
| result = gen( | |
| inpaint_prompt, | |
| source_img, | |
| binary_mask, | |
| new_width, | |
| new_height, | |
| num_steps, | |
| seed, | |
| guidance, | |
| ) | |
| return result, seed | |
| def add_border_and_mask(image, zoom_all=1.0, zoom_left=0, zoom_right=0, zoom_up=0, zoom_down=0, overlap=0.01): | |
| """Adds a black border around the image with individual side control and mask overlap""" | |
| orig_width, orig_height = image.size | |
| # Calculate padding for each side (in pixels) | |
| left_pad = int(orig_width * zoom_left) | |
| right_pad = int(orig_width * zoom_right) | |
| top_pad = int(orig_height * zoom_up) | |
| bottom_pad = int(orig_height * zoom_down) | |
| # Calculate overlap in pixels | |
| overlap_left = int(orig_width * overlap) | |
| overlap_right = int(orig_width * overlap) | |
| overlap_top = int(orig_height * overlap) | |
| overlap_bottom = int(orig_height * overlap) | |
| # If using the all-sides zoom, add it to each side | |
| if zoom_all > 1.0: | |
| extra_each_side = (zoom_all - 1.0) / 2 | |
| left_pad += int(orig_width * extra_each_side) | |
| right_pad += int(orig_width * extra_each_side) | |
| top_pad += int(orig_height * extra_each_side) | |
| bottom_pad += int(orig_height * extra_each_side) | |
| # Calculate new dimensions (ensure they're multiples of 32) | |
| new_width = 32 * round((orig_width + left_pad + right_pad) / 32) | |
| new_height = 32 * round((orig_height + top_pad + bottom_pad) / 32) | |
| # Create new image with black border | |
| bordered_image = Image.new("RGB", (new_width, new_height), (0, 0, 0)) | |
| # Paste original image in position | |
| paste_x = left_pad | |
| paste_y = top_pad | |
| bordered_image.paste(image, (paste_x, paste_y)) | |
| # Create mask (white where the border is, black where the original image was) | |
| mask = Image.new("L", (new_width, new_height), 255) # White background | |
| # Paste black rectangle with overlap adjustment | |
| mask.paste( | |
| 0, | |
| ( | |
| paste_x + overlap_left, # Left edge moves right | |
| paste_y + overlap_top, # Top edge moves down | |
| paste_x + orig_width - overlap_right, # Right edge moves left | |
| paste_y + orig_height - overlap_bottom, # Bottom edge moves up | |
| ), | |
| ) | |
| return bordered_image, mask | |
| def outpaintGen( | |
| img, | |
| outpaint_prompt: str, | |
| overlap: float, | |
| zoom_all: float, | |
| zoom_left: float, | |
| zoom_right: float, | |
| zoom_up: float, | |
| zoom_down: float, | |
| guidance: float, | |
| num_steps: int, | |
| seed: int, | |
| randomize_seed: bool | |
| ): | |
| image = Image.open(img) | |
| new_image, mask_image = add_border_and_mask( | |
| image, | |
| zoom_all=zoom_all, | |
| zoom_left=zoom_left, | |
| zoom_right=zoom_right, | |
| zoom_up=zoom_up, | |
| zoom_down=zoom_down, | |
| overlap=overlap, | |
| ) | |
| width, height = new_image.size | |
| if randomize_seed: | |
| seed = random.randint(0, MAX_SEED) | |
| result = gen( | |
| outpaint_prompt, | |
| new_image, | |
| mask_image, | |
| width, | |
| height, | |
| num_steps, | |
| seed, | |
| guidance, | |
| ) | |
| return result, seed | |
| with gr.Blocks(theme="ocean", title="Flux.1 Fill dev", css=CSS) as demo: | |
| gr.HTML("<h1><center>Flux.1 Fill dev</center></h1>") | |
| gr.HTML(""" | |
| <p> | |
| <center> | |
| FLUX.1 Fill [dev] is a 12 billion parameter rectified flow transformer capable of filling areas in existing images based on a text description. | |
| </center> | |
| </p> | |
| """) | |
| with gr.Tab("Inpainting"): | |
| with gr.Row(): | |
| with gr.Column(): | |
| imgMask = gr.ImageMask(type="filepath", label="Image", layers=False, height=800) | |
| inpaint_prompt = gr.Textbox(label='Prompts ✏️', placeholder="A hat...") | |
| with gr.Row(): | |
| Inpaint_sendBtn = gr.Button(value="Submit", variant='primary') | |
| Inpaint_clearBtn = gr.ClearButton([imgMask, inpaint_prompt], value="Clear") | |
| image_out = gr.Image(type="pil", label="Output", height=960) | |
| with gr.Accordion("Advanced ⚙️", open=False): | |
| guidance = gr.Slider(label="Guidance scale", minimum=1, maximum=50, value=30.0, step=0.1) | |
| num_steps = gr.Slider(label="Steps", minimum=1, maximum=50, value=20, step=1) | |
| seed = gr.Number(label="Seed", value=42, precision=0) | |
| randomize_seed = gr.Checkbox(label="Randomize seed", value=True) | |
| gr.on( | |
| triggers = [ | |
| inpaint_prompt.submit, | |
| Inpaint_sendBtn.click, | |
| ], | |
| fn = inpaintGen, | |
| inputs = [ | |
| imgMask, | |
| inpaint_prompt, | |
| guidance, | |
| num_steps, | |
| seed, | |
| randomize_seed | |
| ], | |
| outputs = [image_out, seed] | |
| ) | |
| with gr.Tab("Outpainting"): | |
| with gr.Row(): | |
| with gr.Column(): | |
| img = gr.Image(type="filepath", label="Image", height=800) | |
| outpaint_prompt = gr.Textbox(label='Prompts ✏️', placeholder="In city...") | |
| with gr.Row(): | |
| outpaint_sendBtn = gr.Button(value="Submit", variant='primary') | |
| outpaint_clearBtn = gr.ClearButton([img, outpaint_prompt], value="Clear") | |
| image_exp = gr.Image(type="pil", label="Output", height=960) | |
| with gr.Accordion("Advanced ⚙️", open=False): | |
| overlap = gr.Slider(label="Overlap", minimum=0.01, maximum=0.25, value=0.01, step=0.01) | |
| zoom_all = gr.Slider(label="Zoom Out Amount (All Sides)", minimum=1.0, maximum=3.0, value=1.0, step=0.1) | |
| with gr.Row(): | |
| zoom_left = gr.Slider(label="Left", minimum=0.0, maximum=1.0, value=0.0, step=0.1) | |
| zoom_right = gr.Slider(label="Right", minimum=0.0, maximum=1.0, value=0.0, step=0.1) | |
| with gr.Row(): | |
| zoom_up = gr.Slider(label="Up", minimum=0.0, maximum=1.0, value=0.0, step=0.1) | |
| zoom_down = gr.Slider(label="Down", minimum=0.0, maximum=1.0, value=0.0, step=0.1) | |
| op_guidance = gr.Slider(label="Guidance scale", minimum=1, maximum=50, value=30.0, step=0.1) | |
| op_num_steps = gr.Slider(label="Steps", minimum=1, maximum=50, value=20, step=1) | |
| op_seed = gr.Number(label="Seed", value=42, precision=0) | |
| op_randomize_seed = gr.Checkbox(label="Randomize seed", value=True) | |
| gr.on( | |
| triggers = [ | |
| outpaint_prompt.submit, | |
| outpaint_sendBtn.click, | |
| ], | |
| fn = outpaintGen, | |
| inputs = [ | |
| img, | |
| outpaint_prompt, | |
| overlap, | |
| zoom_all, | |
| zoom_left, | |
| zoom_right, | |
| zoom_up, | |
| zoom_down, | |
| op_guidance, | |
| op_num_steps, | |
| op_seed, | |
| op_randomize_seed | |
| ], | |
| outputs = [image_exp, op_seed] | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch(show_api=False, share=False) |