Spaces:
Running
on
Zero
Running
on
Zero
| import os | |
| import torch | |
| import spaces | |
| import gradio as gr | |
| from diffusers import FluxFillPipeline | |
| from diffusers.utils import load_image | |
| import random | |
| import numpy as np | |
| from huggingface_hub import hf_hub_download | |
| 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 create_mask_image(mask_array): | |
| # Convert the mask to a numpy array if it's not already | |
| if not isinstance(mask_array, np.ndarray): | |
| mask_array = np.array(mask_array) | |
| # Create a new array with the same shape as the mask, but only for RGB channels | |
| processed_mask = np.zeros((mask_array.shape[0], mask_array.shape[1], 3), dtype=np.uint8) | |
| # Set transparent parts (alpha=0) to black (0, 0, 0) | |
| transparent_mask = mask_array[:, :, 3] == 0 | |
| processed_mask[transparent_mask] = [0, 0, 0] | |
| # Set black parts (RGB=0, 0, 0 and alpha=255) to white (255, 255, 255) | |
| black_mask = (mask_array[:, :, :3] == [0, 0, 0]).all(axis=2) & (mask_array[:, :, 3] == 255) | |
| processed_mask[black_mask] = [255, 255, 255] | |
| return Image.fromarray(processed_mask) | |
| 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] | |
| print(f'source_path: {source_path}') | |
| print(f'mask_path: {mask_path}') | |
| 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 = load_image(source_path) | |
| mask_img = load_image(mask_path) | |
| mask_img = create_mask_image(mask_img) | |
| width, height = source_img.size | |
| if randomize_seed: | |
| seed = random.randint(0, MAX_SEED) | |
| generator = torch.Generator("cpu").manual_seed(seed) | |
| result = pipe( | |
| prompt=inpaint_prompt, | |
| image=source_img, | |
| mask_image=mask_img, | |
| width=width, | |
| height=height, | |
| num_inference_steps=num_steps, | |
| generator=generator, | |
| guidance_scale=guidance, | |
| max_sequence_length=512, | |
| ).images[0] | |
| 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> | |
| A partial redraw of the image based on your prompt words and occluded parts. | |
| </center> | |
| </p> | |
| """) | |
| 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] | |
| ) | |
| if __name__ == "__main__": | |
| demo.queue(api_open=False).launch(show_api=False, share=False) |