Spaces:
Runtime error
Runtime error
| # this code is largely inspired by https://huggingface.co/spaces/hysts/ControlNet-with-Anything-v4/blob/main/app_scribble_interactive.py | |
| # Thank you, hysts! | |
| import sys | |
| sys.path.append('./src/ControlNetInpaint/') | |
| # functionality based on https://github.com/mikonvergence/ControlNetInpaint | |
| import gradio as gr | |
| #import torch | |
| #from torch import autocast // only for GPU | |
| from PIL import Image | |
| import numpy as np | |
| from io import BytesIO | |
| import os | |
| # Usage | |
| # 1. Upload image or fill with white | |
| # 2. Sketch the mask (image->[image,mask] | |
| # 3. Sketch the content of the mask | |
| # Global Storage | |
| CURRENT_IMAGE={'image' : None, | |
| 'mask' : None, | |
| 'guide' : None | |
| } | |
| HEIGHT,WIDTH=512,512 | |
| ## SETUP PIPE | |
| from diffusers import StableDiffusionInpaintPipeline, ControlNetModel, UniPCMultistepScheduler | |
| from src.pipeline_stable_diffusion_controlnet_inpaint import * | |
| from diffusers.utils import load_image | |
| from controlnet_aux import HEDdetector | |
| hed = HEDdetector.from_pretrained('lllyasviel/ControlNet') | |
| controlnet = ControlNetModel.from_pretrained( | |
| "fusing/stable-diffusion-v1-5-controlnet-scribble", torch_dtype=torch.float16 | |
| ) | |
| pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained( | |
| "runwayml/stable-diffusion-inpainting", controlnet=controlnet, torch_dtype=torch.float16 | |
| ) | |
| pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config) | |
| if torch.cuda.is_available(): | |
| # Remove if you do not have xformers installed | |
| # see https://huggingface.co/docs/diffusers/v0.13.0/en/optimization/xformers#installing-xformers | |
| # for installation instructions | |
| pipe.enable_xformers_memory_efficient_attention() | |
| pipe.to('cuda') | |
| # Functions | |
| def get_guide(image): | |
| return hed(image,scribble=True) | |
| def create_demo(max_images=12, default_num_images=3): | |
| with gr.Blocks(theme=gr.themes.Default(font=[gr.themes.GoogleFont("IBM Plex Mono"), "ui-monospace","monospace"])) as demo: | |
| gr.Markdown('# Cut and Sketch ✂️▶️✏️') | |
| with gr.Accordion('Instructions', open=False): | |
| gr.Markdown('## Cut ✂️') | |
| gr.Markdown('1. Upload your image below') | |
| gr.Markdown('2. **Draw the mask** for the region you want changed (Cut ✂️)') | |
| gr.Markdown('3. Click `Set Mask` when it is ready!') | |
| gr.Markdown('## Sketch ✏️') | |
| gr.Markdown('4. Now, you can **sketch a replacement** object! (Sketch ✏️)') | |
| gr.Markdown('5. (You can also provide a **text prompt** if you want)') | |
| gr.Markdown('6. 🔮 Click `Generate` when ready! ') | |
| with gr.Row() as main_blocks: | |
| with gr.Column() as step_1: | |
| gr.Markdown('### Mask Input') | |
| input_image = gr.Image(source='upload', | |
| shape=[HEIGHT,WIDTH], | |
| type='numpy', | |
| elem_id="image-upload", | |
| label='Mask Draw (Cut!)', | |
| tool='sketch', | |
| brush_radius=80) | |
| mask_button = gr.Button(label='Set Mask', value='Set Mask',full_width=False) | |
| with gr.Column(visible=False) as step_2: | |
| gr.Markdown('### Sketch Input') | |
| sketch_image = gr.Image(source='upload', | |
| shape=[HEIGHT,WIDTH], | |
| type='numpy', | |
| label='Fill Draw (Sketch!)', | |
| tool='sketch', | |
| brush_radius=15) | |
| run_button = gr.Button(label='Generate', value='Generate') | |
| prompt = gr.Textbox(label='Prompt') | |
| with gr.Column() as output_step: | |
| gr.Markdown('### Output') | |
| output_image = gr.Gallery( | |
| label="Generated images", | |
| show_label=False, | |
| elem_id="gallery", | |
| ) | |
| with gr.Accordion('Advanced options', open=False): | |
| num_steps = gr.Slider(label='Steps', | |
| minimum=1, | |
| maximum=100, | |
| value=20, | |
| step=1) | |
| text_scale = gr.Slider(label='Text Guidance Scale', | |
| minimum=0.1, | |
| maximum=30.0, | |
| value=7.5, | |
| step=0.1) | |
| seed = gr.Slider(label='Seed', | |
| minimum=-1, | |
| maximum=2147483647, | |
| step=1, | |
| randomize=True) | |
| sketch_scale = gr.Slider(label='Sketch Guidance Scale', | |
| minimum=0.0, | |
| maximum=1.0, | |
| value=1.0, | |
| step=0.05) | |
| inputs = [ | |
| sketch_image, | |
| prompt, | |
| num_steps, | |
| text_scale, | |
| sketch_scale, | |
| seed | |
| ] | |
| # STEP 1: Set Mask | |
| def set_mask(image): | |
| img=image['image'][...,:3] | |
| mask=1*(image['mask'][...,:3]>0) | |
| # save vars | |
| CURRENT_IMAGE['image']=img | |
| CURRENT_IMAGE['mask']=mask | |
| guide=get_guide(img) | |
| CURRENT_IMAGE['guide']=np.array(guide) | |
| guide=255-np.asarray(guide) | |
| seg_img = guide*(1-mask) + mask*192 | |
| preview = img * (seg_img==255) | |
| vis_image=(preview/2).astype(seg_img.dtype) + seg_img * (seg_img!=255) | |
| return {input_image : image['image'], | |
| sketch_image : vis_image, | |
| step_1: gr.update(visible=False), | |
| step_2: gr.update(visible=True) | |
| } | |
| # STEP 2: Generate | |
| def generate(image, | |
| prompt, | |
| num_steps, | |
| text_scale, | |
| sketch_scale, | |
| seed): | |
| sketch=(255*(image['mask'][...,:3]>0)).astype(CURRENT_IMAGE['image'].dtype) | |
| mask=CURRENT_IMAGE['mask'] | |
| CURRENT_IMAGE['guide']=(CURRENT_IMAGE['guide']*(mask==0) + sketch*(mask!=0)).astype(CURRENT_IMAGE['image'].dtype) | |
| mask_img=255*CURRENT_IMAGE['mask'].astype(CURRENT_IMAGE['image'].dtype) | |
| new_image = pipe( | |
| prompt, | |
| num_inference_steps=num_steps, | |
| guidance_scale=text_scale, | |
| generator=torch.manual_seed(seed), | |
| image=Image.fromarray(CURRENT_IMAGE['image']), | |
| control_image=Image.fromarray(CURRENT_IMAGE['guide']), | |
| controlnet_conditioning_scale=sketch_scale, | |
| mask_image=Image.fromarray(mask_img) | |
| ).images#[0] | |
| return {output_image : new_image, | |
| step_1: gr.update(visible=True), | |
| step_2: gr.update(visible=False) | |
| } | |
| mask_button.click(fn=set_mask, inputs=[input_image], outputs=[input_image, sketch_image, step_1,step_2]) | |
| run_button.click(fn=generate, inputs=inputs, outputs=[output_image, step_1,step_2]) | |
| return demo | |
| if __name__ == '__main__': | |
| demo = create_demo() | |
| demo.queue().launch() |