Spaces:
Running
Running
import gradio as gr | |
from PIL import Image, ImageDraw | |
def run_afm_app(task_selector, input_image, mask_image, text_input, text_input_x, text_input_gsam, coord_input, | |
ddim_steps, ddim_steps_pipe, inpaint_input_gsam, text_input_inpaint_pipe, text_input_restyling, | |
blur, sharpen, prompt_outpaint, e_l, e_r, e_u, e_d, steps_outpaint, prompt_background , steps_br, | |
str_res, gs_res, np_res, steps_res, np_inpaint, steps_inpaint, prompt_txt2img, np_txt2img, gs_txt2img, | |
steps_txt2img, steps_super, dilation_bool, dilation_value, steps_inp): | |
print(f"Task selected: {task_selector}") | |
if task_selector == "SAM": | |
from mask_sam import sam_gradio | |
return sam_gradio(input_image, coord_input, dilation_bool, dilation_value) | |
if task_selector == "GroundedSAM": | |
from mask_groundedsam import groundedsam_mask_gradio | |
return groundedsam_mask_gradio(input_image, text_input, dilation_bool, dilation_value) | |
if task_selector == "Stable Diffusion with ControlNet Inpainting": | |
from inpaint_sd_controlnet import controlnet_inpaint_gradio | |
return controlnet_inpaint_gradio(input_image, mask_image, text_input_x) | |
if task_selector == "Stable Diffusion v1.5 Inpainting": | |
from inpaint_sd import inpaint_sd_gradio | |
return inpaint_sd_gradio(input_image, mask_image, text_input_x, steps_inp) | |
if task_selector == "Stable Diffusion XL Inpainting": | |
from inpaint_sdxl import inpaint_sdxl_gradio | |
return inpaint_sdxl_gradio(input_image, mask_image, text_input_x, steps_inp) | |
if task_selector == "Kandinsky v2.2 Inpainting": | |
from inpaint_kandinsky import inpaint_kandinsky_gradio | |
return inpaint_kandinsky_gradio(input_image, mask_image, text_input_x, steps_inp) | |
if task_selector == "GroundedSAM Inpainting": | |
from inpaint_groundedsam import groundedsam_inpaint_gradio | |
return groundedsam_inpaint_gradio(input_image, text_input_gsam, inpaint_input_gsam) | |
if task_selector == "Object Removal LDM": | |
from eraser_ldm import ldm_removal_gradio | |
return ldm_removal_gradio(input_image, mask_image, ddim_steps) | |
if task_selector == "Restyling - Stable Diffusion v1.5": | |
from restyling_sd import restyling_gradio | |
return restyling_gradio(input_image, text_input_restyling, str_res, gs_res, np_res, steps_res) | |
if task_selector == "Restyling - Stable Diffusion XL": | |
from restyling_sdxl import restyling_sdxl_gradio | |
return restyling_sdxl_gradio(input_image, text_input_restyling, str_res, gs_res, np_res, steps_res) | |
if task_selector == "Restyling - Kandinsky v2.2": | |
from restyling_kandinsky import restyling_kandinsky_gradio | |
return restyling_kandinsky_gradio(input_image, text_input_restyling, str_res, gs_res, np_res, steps_res) | |
if task_selector == "Superresolution - LDM x4 OpenImages": | |
from superres_ldm import superres_gradio | |
return superres_gradio(input_image, steps_super) | |
if task_selector == "Superresolution - Stability AI x4 Upscaler": | |
from superres_upscaler import superres_upscaler_gradio | |
return superres_upscaler_gradio(input_image, steps_super) | |
if task_selector == "LDM Removal Pipeline": | |
from eraser_ldm_pipe import ldm_removal_pipe_gradio | |
return ldm_removal_pipe_gradio(input_image, coord_input, ddim_steps_pipe) | |
if task_selector in ["Stable Diffusion v1.5 Inpainting Pipeline", "Stable Diffusion XL Inpainting Pipeline", "Kandinsky v2.2 Inpainting Pipeline"]: | |
from inpaint_pipe import inpaint_pipe_gradio | |
return inpaint_pipe_gradio(task_selector, input_image, coord_input, text_input_inpaint_pipe, np_inpaint, steps_inpaint) | |
if task_selector == "Stable Diffusion with ControlNet Inpainting Pipeline": | |
from inpaint_sd_controlnet_pipe import inpaint_func_pipe_gradio | |
return inpaint_func_pipe_gradio(input_image, coord_input, text_input_inpaint_pipe, np_inpaint, steps_inpaint) | |
if task_selector == "Portrait Mode - Depth Anything": | |
from blur_image import portrait_gradio | |
return portrait_gradio(input_image, blur, sharpen) | |
if task_selector == "Outpainting - Stable Diffusion": | |
from outpaint_sd import outpaint_sd_gradio | |
return outpaint_sd_gradio(input_image, prompt_outpaint, e_l, e_r, e_u, e_d, steps_outpaint) | |
if task_selector == "Outpainting - Stable Diffusion XL": | |
from outpaint_sdxl import outpaint_sdxl_gradio | |
return outpaint_sdxl_gradio(input_image, prompt_outpaint, e_l, e_r, e_u, e_d, steps_outpaint) | |
if task_selector == "Background Replacement - Stable Diffusion": | |
from background_replace_sd import background_replace_sd_gradio | |
return background_replace_sd_gradio(input_image, prompt_background , steps_br) | |
if task_selector == "Background Replacement - Stable Diffusion XL": | |
from background_replace_sdxl import background_replace_sdxl_gradio | |
return background_replace_sdxl_gradio(input_image, prompt_background , steps_br) | |
if task_selector in ["Stable Diffusion v1.5 Txt2Img", "Stable Diffusion XL Txt2Img", "Kandinsky v2.2 Txt2Img"]: | |
from txt2img_generation import txt2img_gradio | |
return txt2img_gradio(input_image, task_selector, prompt_txt2img, np_txt2img, gs_txt2img, steps_txt2img) | |
if task_selector == "Eraser - LaMa": | |
from eraser_lama import eraser_lama_gradio | |
return eraser_lama_gradio(input_image, mask_image) | |
selected_points = [] | |
def input_handler(evt: gr.SelectData, input_image): | |
global selected_points | |
coords = evt.index | |
x, y = coords[0], coords[1] | |
selected_points.append([x, y]) | |
coord_string = '; '.join([f"{pt[0]},{pt[1]}" for pt in selected_points]) | |
image_with_points = input_image.copy() | |
draw = ImageDraw.Draw(image_with_points) | |
for point in selected_points: | |
draw.ellipse((point[0] - 2, point[1] - 2, point[0] + 2, point[1] + 2), fill="red", outline="red") | |
return coord_string, image_with_points | |
def reset_selected_points(input_image): | |
global selected_points | |
selected_points = [] | |
print("Selected points have been reset.") | |
return "", input_image | |
def reload_image(original_image_path): | |
original_image = original_image_path | |
return original_image | |
def update_task_selector(task_selector, task): | |
return task | |
def reload_image_with_output(output_image): | |
return output_image | |
def reload_mask(output_image): | |
return output_image | |
title = "# AFM Image-Editing App" | |
if __name__ == "__main__": | |
block = gr.Blocks(theme='shivi/calm_seafoam') | |
with block: | |
gr.Markdown(title) | |
gr.Markdown( | |
""" | |
WARNING: To run this app on GPU, please follow the instructions on the repo: https://github.com/matt576/image-editing | |
Welcome to the AFM Image-Editing App! | |
First, upload an input image or generate it via Txt2Img below. | |
Then, choose the desired task by navigating the tabs. | |
Finally, choose the model on the Dropdown within each tab and click on 'Generate'! Enjoy the App! | |
""") | |
original_image_path = "scott.png" # Select input image path here | |
# original_image_path = "outputs/txt2img/generated_input.png" # for txt2img generated input image | |
input_mask_path = "scott.png" # Optional, make sure it matches the input image | |
original_image = Image.open(original_image_path) | |
with gr.Row(): | |
with gr.Column(): | |
input_image = gr.Image(label="Input Image", sources='upload', type="pil", value=original_image_path, interactive=True) | |
with gr.Column(): | |
output_image = gr.Image(label="Generated Image", type="pil") | |
with gr.Row(): | |
generate_button = gr.Button("Generate!") | |
with gr.Row(): | |
with gr.Column(): | |
gr.Markdown("Type image coordinates manually or click on the image directly:") | |
coord_input = gr.Textbox(label="Pixel Coordinates (x,y), Format x1,y1; x2,y2 ...", value="") | |
reset_button = gr.Button("Reset coordinates") | |
reload_image_button = gr.Button("Clear Image") | |
reload_output_button = gr.Button("Load Output") | |
task_selector = gr.State(value="") | |
with gr.Accordion("Txt2Img Generation (Optional)", open=False): | |
tab_task_selector_11 = gr.Dropdown(["Stable Diffusion v1.5 Txt2Img", | |
"Stable Diffusion XL Txt2Img", | |
"Kandinsky v2.2 Txt2Img"], label="Select Model") | |
gr.Markdown(""" | |
### Instructions | |
Use this feature if you wish to generate your own input image. | |
After generation, simply uncomment the original_image_path line on the gradio script and relaunch the app! | |
Required Inputs: Text Prompt, str_res, gs_res, np_res, steps_res | |
Example prompt: "Photorealistic Gotham City night skyline, rain pouring down, dark clouds with streaks of lightning." | |
Example negative prompt: "poor details, poor quality, blurry, deformed, extra limbs" | |
""") | |
prompt_txt2img = gr.Textbox(label="Text Prompt: ", value="Photorealistic Gotham City night skyline, Batman standing on top of skyscraper, close shot, unreal engine, cinematic, rain pouring down, dark clouds with streaks of lightning") | |
np_txt2img = gr.Textbox(label="Negative Prompt", value="poor details, poor quality, blurry, deformed, extra limbs") | |
gs_txt2img = gr.Slider(minimum=0.0, maximum=50.0, label="Guidance Scale", value=7.5) | |
steps_txt2img = gr.Slider(minimum=5, maximum=200, label="Number of inference steps", value=30, step=1) | |
with gr.Accordion("Mask Input Tasks (Optional)", open=False): | |
gr.Markdown(""" | |
Here is the mask uploaded directly from the gradio script, if you wish to change it, | |
use the Mask Generation Preview Tab and click the 'Load Preview Mask' button. | |
""") | |
mask_image = gr.Image(label="Input Mask (Optional)", sources='upload', type="pil", value=input_mask_path) | |
with gr.Tab("Inpainting - Object Replacement"): | |
tab_task_selector_2 = gr.Dropdown(["Stable Diffusion with ControlNet Inpainting", | |
"Stable Diffusion v1.5 Inpainting", | |
"Stable Diffusion XL Inpainting", | |
"Kandinsky v2.2 Inpainting"], | |
label="Select Model") | |
gr.Markdown(""" | |
### Instructions | |
All models in this section work with the given uploaded input mask. | |
Required Inputs: Input Mask (Upload) , Text Prompt - Object to replace masked area on given input mask below. | |
Input in the text box below the desired object to be inpainted in place of the mask input below. | |
Example prompt: "astronaut, white suit, 8k, extremely detailed, ornate, cinematic lighting, vivid, photorealistic, detailed, high quality" | |
""") | |
text_input_x = gr.Textbox(label="Text Prompt: ", value="astronaut, white suit, 8k, extremely detailed, ornate, cinematic lighting, vivid, photorealistic, detailed, high quality") | |
steps_inp = gr.Slider(minimum=5, maximum=200, label="Number of inference steps: ", value=50, step=1) | |
with gr.Tab("Object Removal"): | |
tab_task_selector_3 = gr.Dropdown(["Object Removal LDM", "Eraser - LaMa"], label="Select Model") | |
gr.Markdown(""" | |
### Instructions | |
- **Object Removal LDM**: | |
Required inputs: Input image, Input Mask (Upload or from Preview), DDIM Steps | |
Given the uploaded mask below, simply adjust the slider below according to the desired number of iterations. | |
- **Eraser - LaMa**: | |
Required inputs: Input image, Input Mask (Upload or from Preview) | |
Please note, due to compability issues with the LaMa model and our gradio app, the output visualiztion will not | |
work in the app, but your output will be saved to: code/outputs/untracked/eraser-lama. | |
""") | |
ddim_steps = gr.Slider(minimum=5, maximum=250, label="Number of DDIM sampling steps for object removal LDM", value=150, step=1) | |
with gr.Column(): | |
with gr.Tab("Mask Generation Preview"): | |
tab_task_selector_1 = gr.Dropdown(["SAM", "GroundedSAM"], label="Select Model") | |
reload_mask_button = gr.Button("Load Preview Mask") | |
gr.Markdown(""" | |
### Instructions | |
- **SAM**: | |
Required inputs: Input Image, Pixel Coordinates, (Optional) Dilation | |
Type image coordinates manually or click on the image directly. Finally, simply click on the 'Generate' button. | |
""") | |
dilation_bool = gr.Dropdown(["Yes", "No"], label="Use dilation (recommended for inpainting)") | |
dilation_value = gr.Slider(minimum=0, maximum=50, label="Dilation value (recommended: 10) ", value=10, step = 1) | |
gr.Markdown(""" | |
- **GroundedSAM (GroundingDINO + SAM)**: | |
Required Inputs: Text Prompt [object(s) to be detected], (Optional) Dilation | |
Input in the text box below the object(s) in the input image for which the masks are to be generated. | |
""") | |
text_input = gr.Textbox(label="Text Prompt: ", value="dog") | |
with gr.Tab("Restyling"): | |
tab_task_selector_4 = gr.Dropdown(["Restyling - Stable Diffusion v1.5", | |
"Restyling - Stable Diffusion XL", | |
"Restyling - Kandinsky v2.2"], label="Select Model") | |
gr.Markdown(""" | |
### Instructions | |
Required Inputs: Input Image, Text Prompt, str_res, gs_res, np_res, steps_res | |
Example Text Prompt: "Photorealistic Gotham City night skyline, rain pouring down, dark clouds with streaks of lightning." | |
Example Negative Prompt: "poor details, poor quality, blurry, deformed, extra limbs" | |
""") | |
text_input_restyling = gr.Textbox(label="Text Prompt: ", value="Futuristic night city from Cyberpunk 2077, rainy night, close shot, 35 mm, realism, octane render, 8 k, exploration, cinematic, pixbay, modernist, realistic, unreal engine, hyper detailed, photorealistic, maximum detail, volumetric light, moody cinematic epic concept art, vivid") | |
str_res = gr.Slider(minimum=0.1, maximum=1.0, label="Strength: ", value=0.75, step=0.01) | |
gs_res = gr.Slider(minimum=0.0, maximum=50.0, label="Guidance Scale: ", value=7.5, step=0.1) | |
np_res = gr.Textbox(label="Negative Prompt: ", value="poor details, poor quality, blurry, deformed, extra limbs") | |
steps_res = gr.Slider(minimum=5, maximum=150, label="Number of inference steps: ", value=30, step=1) | |
with gr.Tab("Superresolution"): | |
tab_task_selector_5 = gr.Dropdown(["Superresolution - LDM x4 OpenImages", | |
"Superresolution - Stability AI x4 Upscaler"], label="Select Model") | |
gr.Markdown(""" | |
### Instructions | |
Required Inputs: Input Image, Number of Inference Steps | |
Select model on the Dropdown menu, number of inference steps, and click the 'Generate' button to get your new image. | |
""") | |
steps_super = gr.Slider(minimum=5, maximum=150, label="Number of inference steps: ", value=30, step=1) | |
with gr.Tab("Pipeline: Inpainting - Object Replacement"): | |
tab_task_selector_6 = gr.Dropdown(["GroundedSAM Inpainting", | |
"Stable Diffusion with ControlNet Inpainting Pipeline", | |
"Stable Diffusion v1.5 Inpainting Pipeline", | |
"Stable Diffusion XL Inpainting Pipeline", | |
"Kandinsky v2.2 Inpainting Pipeline"], label="Select Model") | |
gr.Markdown(""" | |
- **GroundedSAM Inpainting (GroundingDINO + SAM + Stable Diffusion)**: | |
Required Inputs: Input Image, Detection Prompt , Inpainting Prompt | |
Input in the text box below the object(s) in the input image for which the masks are to be generated. | |
Example detection prompt: "dog" | |
Example inpaint prompt: "white tiger, photorealistic, detailed, high quality" | |
""") | |
text_input_gsam = gr.Textbox(label="Detection Prompt: ", value="dog") | |
inpaint_input_gsam = gr.Textbox(label="Inpainting Prompt: ", value="astronaut, white suit, 8k, extremely detailed, ornate, cinematic lighting, vivid, photorealistic, detailed, high quality") | |
gr.Markdown(""" | |
- **Kandinsky v2.2 / Stable Diffusion v1.5 / SDXL / SD + ControlNet**: | |
Required Inputs: Input Image, Pixel Coodinates , Inpainting Prompt | |
Input in the text box below the object(s) in the input image for which the masks are to be generated. | |
Example Text Prompt: "white tiger, photorealistic, detailed, high quality" | |
Example Negative Prompt: "poor details, poor quality, blurry, deformed, extra limbs" | |
""") | |
text_input_inpaint_pipe = gr.Textbox(label="Text Prompt: ", value="astronaut, white suit, 8k, extremely detailed, ornate, cinematic lighting, vivid, photorealistic, detailed, high quality") | |
np_inpaint = gr.Textbox(label="Negative Prompt: ", value="poor details, poor quality, blurry, deformed, extra limbs") | |
steps_inpaint = gr.Slider(minimum=5, maximum=200, label="Number of inference steps: ", value=150, step=1) | |
with gr.Tab("Pipeline - Object Removal"): | |
tab_task_selector_7 = gr.Dropdown(["LDM Removal Pipeline", " "], label="Select Model") | |
gr.Markdown(""" | |
### Instructions | |
- **LDM Removal Pipeline**: | |
Required inputs: Input Image, Pixel Coodinates, DDIM Steps | |
If you wish to view the mask before the fnal output, go to the 'Mask Generation Preview' Tab. | |
Type the image coordinates manually in the box under the image or click on the image directly. | |
For a more detailed mask of a specific object or part of it, select multiple points. | |
Finally, choose number of DDIM steps simply click on the 'Generate' button: | |
""") | |
ddim_steps_pipe = gr.Slider(minimum=5, maximum=250, label="Number of DDIM sampling steps for object removal", value=150, step=1) | |
with gr.Tab("Background Blurring"): | |
tab_task_selector_8 = gr.Dropdown(["Portrait Mode - Depth Anything"], label='Select Model') | |
gr.Markdown(""" | |
### Instructions | |
- **Portrait Mode - Depth Anything**: | |
Required inputs: Input Image, box blur, sharpen | |
Recommended blur values range: 2-25 | |
Recommended sharpen values range: 0-5 | |
Adjust the required inputs with the siders below: | |
""") | |
blur = gr.Slider(minimum=0, maximum=50, label="Box Blur value", value=5, step=1) | |
sharpen = gr.Slider(minimum=0, maximum=7, label="Sharpen Parameter", value=0, step=1) | |
with gr.Tab("Outpainting"): | |
tab_task_selector_9 = gr.Dropdown(["Outpainting - Stable Diffusion", "Outpainting - Stable Diffusion XL"], label='Select Model') | |
gr.Markdown(""" | |
### Instructions | |
- **Outpainting - Stable Diffusion**: | |
Required inputs: Input Image, Text Prompt, extend left/right/up/down, steps | |
Choose how much and which direction you want to extend /outpaint your image and specify a text prompt. | |
Example prompt: "open plan, kitchen and living room, black umbrella on the floor, modular furniture with cotton textiles, wooden floor, high ceiling, large steel windows viewing a city" | |
""") | |
prompt_outpaint = gr.Textbox(label="Text Prompt: ", value="open plan, kitchen and living room, black umbrella on the floor, modular furniture with cotton textiles, wooden floor, high ceiling, large steel windows viewing a city") | |
e_l = gr.Slider(minimum=0, maximum=1000, label="Extend Left", value=200, step=1) | |
e_r = gr.Slider(minimum=0, maximum=1000, label="Extend Right", value=200, step=1) | |
e_u = gr.Slider(minimum=0, maximum=1000, label="Extend Up", value=200, step=1) | |
e_d = gr.Slider(minimum=0, maximum=1000, label="Extend Down", value=200, step=1) | |
steps_outpaint = gr.Slider(minimum=0, maximum=200, label="Number of Steps", value=50, step=1) | |
with gr.Tab("Background Replacement"): | |
tab_task_selector_10 = gr.Dropdown(["Background Replacement - Stable Diffusion", "Background Replacement - Stable Diffusion XL"], label='Select Model') | |
gr.Markdown(""" | |
### Instructions | |
- **Background Replacement - Stable Diffusion**: | |
Required inputs: Input Image, Text Prompt, steps | |
Specify the new background in the text box below. | |
Example prompt: "dog sitting on the beach, sunny day, blue sky" | |
""") | |
prompt_background = gr.Textbox(label="Text Prompt: ", value="dog sitting on the beach, sunny day, blue sky, cinematic, pixbay, modernist, realistic, unreal engine, hyper detailed, photorealistic, maximum detail, volumetric light, moody cinematic epic concept art, vivid") | |
steps_br = gr.Slider(minimum=0, maximum=200, label="Number of Steps", value=30, step=1) | |
input_image.select(input_handler, inputs=[input_image], outputs=[coord_input, input_image]) | |
generate_button.click( | |
fn=run_afm_app, | |
inputs=[task_selector, input_image, mask_image, text_input, text_input_x, text_input_gsam, coord_input, ddim_steps, ddim_steps_pipe, | |
inpaint_input_gsam, text_input_inpaint_pipe, text_input_restyling, blur, sharpen, prompt_outpaint, e_l, e_r, e_u, e_d, steps_outpaint, | |
prompt_background, steps_br, str_res, gs_res, np_res, steps_res, np_inpaint, steps_inpaint, prompt_txt2img, np_txt2img, gs_txt2img, | |
steps_txt2img, steps_super, dilation_bool, dilation_value, steps_inp], | |
outputs=output_image | |
) | |
reset_button.click( | |
fn=reset_selected_points, | |
inputs=[input_image], | |
outputs=[coord_input, input_image] | |
) | |
reload_image_button.click( | |
fn=reload_image, | |
inputs=[gr.State(original_image_path)], | |
outputs=[input_image] | |
) | |
reload_output_button.click( | |
fn=reload_image_with_output, | |
inputs=[output_image], | |
outputs=[input_image] | |
) | |
reload_mask_button.click( | |
fn=reload_mask, | |
inputs=[output_image], | |
outputs=[mask_image] | |
) | |
tab_task_selector_1.change(fn=update_task_selector, inputs=[task_selector, tab_task_selector_1], outputs=[task_selector]) | |
tab_task_selector_2.change(fn=update_task_selector, inputs=[task_selector, tab_task_selector_2], outputs=[task_selector]) | |
tab_task_selector_3.change(fn=update_task_selector, inputs=[task_selector, tab_task_selector_3], outputs=[task_selector]) | |
tab_task_selector_4.change(fn=update_task_selector, inputs=[task_selector, tab_task_selector_4], outputs=[task_selector]) | |
tab_task_selector_5.change(fn=update_task_selector, inputs=[task_selector, tab_task_selector_5], outputs=[task_selector]) | |
tab_task_selector_6.change(fn=update_task_selector, inputs=[task_selector, tab_task_selector_6], outputs=[task_selector]) | |
tab_task_selector_7.change(fn=update_task_selector, inputs=[task_selector, tab_task_selector_7], outputs=[task_selector]) | |
tab_task_selector_8.change(fn=update_task_selector, inputs=[task_selector, tab_task_selector_8], outputs=[task_selector]) | |
tab_task_selector_9.change(fn=update_task_selector, inputs=[task_selector, tab_task_selector_9], outputs=[task_selector]) | |
tab_task_selector_10.change(fn=update_task_selector, inputs=[task_selector, tab_task_selector_10], outputs=[task_selector]) | |
tab_task_selector_11.change(fn=update_task_selector, inputs=[task_selector, tab_task_selector_11], outputs=[task_selector]) | |
block.launch(share=True) |