Spaces:
Runtime error
Runtime error
| from PIL import Image | |
| import streamlit as st | |
| from streamlit_drawable_canvas import st_canvas | |
| from streamlit_lottie import st_lottie | |
| from streamlit_option_menu import option_menu | |
| import requests | |
| import os | |
| import cv2 | |
| import einops | |
| import gradio as gr | |
| import numpy as np | |
| import torch | |
| import random | |
| from huggingface_hub import hf_hub_download | |
| from pytorch_lightning import seed_everything | |
| from annotator.util import resize_image, HWC3 | |
| from annotator.hed import HEDdetector, nms | |
| from cldm.model import create_model, load_state_dict | |
| from cldm.ddim_hacked import DDIMSampler | |
| st.set_page_config( | |
| page_title="ControllNet", | |
| page_icon="🖥️", | |
| layout="wide", | |
| initial_sidebar_state="expanded" | |
| ) | |
| save_memory = False | |
| def load_model(): | |
| model_path = hf_hub_download('lllyasviel/ControlNet', 'models/control_sd15_scribble.pth') | |
| model = create_model('./models/cldm_v15.yaml').cpu() | |
| if torch.cuda.is_available(): | |
| model.load_state_dict(load_state_dict(model_path, location='cuda')) | |
| model = model.cuda() | |
| else: | |
| model.load_state_dict(load_state_dict(model_path, location='cpu')) | |
| return model | |
| def process(input_image, prompt, a_prompt, n_prompt, num_samples, image_resolution, detect_resolution, ddim_steps, guess_mode, strength, scale, seed, eta): | |
| with torch.no_grad(): | |
| input_image = HWC3(input_image[:, :, 0]) | |
| detected_map = apply_hed(resize_image(input_image, detect_resolution)) | |
| detected_map = HWC3(detected_map) | |
| img = resize_image(input_image, image_resolution) | |
| H, W, C = img.shape | |
| detected_map = cv2.resize(detected_map, (W, H), interpolation=cv2.INTER_LINEAR) | |
| detected_map = nms(detected_map, 127, 3.0) | |
| detected_map = cv2.GaussianBlur(detected_map, (0, 0), 3.0) | |
| detected_map[detected_map > 4] = 255 | |
| detected_map[detected_map < 255] = 0 | |
| if torch.cuda.is_available(): | |
| control = torch.from_numpy(detected_map.copy()).float().cuda() / 255.0 | |
| else: | |
| control = torch.from_numpy(detected_map.copy()).float() / 255.0 | |
| control = torch.stack([control for _ in range(num_samples)], dim=0) | |
| control = einops.rearrange(control, 'b h w c -> b c h w').clone() | |
| if seed == -1: | |
| seed = random.randint(0, 2147483647) | |
| seed_everything(seed) | |
| if save_memory: | |
| model.low_vram_shift(is_diffusing=False) | |
| cond = {"c_concat": [control], "c_crossattn": [model.get_learned_conditioning([prompt + ', ' + a_prompt] * num_samples)]} | |
| un_cond = {"c_concat": None if guess_mode else [control], "c_crossattn": [model.get_learned_conditioning([n_prompt] * num_samples)]} | |
| shape = (4, H // 8, W // 8) | |
| if save_memory: | |
| model.low_vram_shift(is_diffusing=True) | |
| model.control_scales = [strength * (0.825 ** float(12 - i)) for i in range(13)] if guess_mode else ([strength] * 13) # Magic number. IDK why. Perhaps because 0.825**12<0.01 but 0.826**12>0.01 | |
| samples, intermediates = ddim_sampler.sample(ddim_steps, num_samples, | |
| shape, cond, verbose=False, eta=eta, | |
| unconditional_guidance_scale=scale, | |
| unconditional_conditioning=un_cond) | |
| if save_memory: | |
| model.low_vram_shift(is_diffusing=False) | |
| x_samples = model.decode_first_stage(samples) | |
| x_samples = (einops.rearrange(x_samples, 'b c h w -> b h w c') * 127.5 + 127.5).cpu().numpy().clip(0, 255).astype(np.uint8) | |
| results = [x_samples[i] for i in range(num_samples)] | |
| # return [255 - detected_map] + results | |
| return results | |
| def load_lottieurl(url: str): | |
| r = requests.get(url) | |
| if r.status_code != 200: | |
| return None | |
| return r.json() | |
| model = load_model() | |
| ddim_sampler = DDIMSampler(model) | |
| apply_hed = HEDdetector() | |
| def main(): | |
| lottie_penguin = load_lottieurl('https://assets5.lottiefiles.com/datafiles/B8q1AyJ5t1wb5S8a2ggTqYNxS1WiKN9mjS76TBpw/articulation/articulation.json') | |
| st.header('Draw and generate image with ControlNet') | |
| with st.sidebar: | |
| st_lottie(lottie_penguin, height=200) | |
| choose = option_menu("Generate image", ["Canvas", "Upload"], | |
| icons=['file-plus', 'cloud-upload'], | |
| menu_icon="infinity", default_index=0, | |
| styles={ | |
| "container": {"padding": ".0rem", "font-size": "14px"}, | |
| "nav-link-selected": {"color": "#000000", "font-size": "16px"}, | |
| } | |
| ) | |
| st.sidebar.markdown( | |
| """ | |
| ___ | |
| <p style='text-align: center'> | |
| ControlNet is as fast as fine-tuning a diffusion model to support additional input conditions | |
| <br/> | |
| <a href="https://arxiv.org/abs/2302.05543" target="_blank">Article</a> | |
| </p> | |
| <p style='text-align: center; font-size: 14px;'> | |
| Spaces creating by | |
| <br/> | |
| <a href="https://www.linkedin.com/in/vumichien/" target="_blank">Chien Vu</a> | |
| <br/> | |
| <img src='https://visitor-badge.glitch.me/badge?page_id=Canvas.ControlNet' alt='visitor badge'> | |
| </p> | |
| """, | |
| unsafe_allow_html=True, | |
| ) | |
| if choose == 'Upload': | |
| st.info("Upload your own scribbles, fill the prompt and enjoy") | |
| with st.form(key='generate_form'): | |
| upload_file = st.file_uploader("Upload image", type=["png", "jpg", "jpeg"]) | |
| prompt = st.text_input(label="Prompt", placeholder='Type your instruction') | |
| col11, col12 = st.columns(2) | |
| with st.expander('Advanced option', expanded=False): | |
| col21, col22 = st.columns(2) | |
| with col21: | |
| image_resolution = st.slider(label="Image Resolution", min_value=256, max_value=512, value=512, step=256) | |
| strength = st.slider(label="Control Strength", min_value=0.0, max_value=2.0, value=1.0, step=0.01) | |
| guess_mode = st.checkbox(label='Guess Mode', value=False) | |
| detect_resolution = st.slider(label="HED Resolution", min_value=128, max_value=1024, value=512, step=1) | |
| ddim_steps = st.slider(label="Steps", min_value=1, max_value=100, value=20, step=1) | |
| with col22: | |
| scale = st.slider(label="Guidance Scale", min_value=0.1, max_value=30.0, value=9.0, step=0.1) | |
| seed = st.number_input(label="Seed", min_value=-1, value=-1) | |
| eta = st.number_input(label="eta (DDIM)", value=0.0) | |
| a_prompt = st.text_input(label="Added Prompt", value='best quality, extremely detailed') | |
| n_prompt = st.text_input(label="Negative Prompt", | |
| value='longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality') | |
| generate_button = st.form_submit_button(label='Generate Image') | |
| if upload_file: | |
| input_image = np.asarray(Image.open(upload_file).convert("RGB")) | |
| print("input_image", input_image.shape) | |
| if generate_button: | |
| with st.spinner(text=f"It may take up to 1 minute under high load. Generating images..."): | |
| results = process(input_image, prompt, a_prompt, n_prompt, 1, image_resolution, detect_resolution, ddim_steps, guess_mode, strength, scale, seed, eta) | |
| print("input_image", input_image.shape) | |
| print("results", results[0].shape) | |
| H, W, C = input_image.shape | |
| output_image = cv2.resize(results[0], (W, H), interpolation=cv2.INTER_AREA) | |
| col11.image(input_image, channels='RGB', width=None, clamp=False, caption='Input image') | |
| col12.image(output_image, channels='RGB', width=None, clamp=False, caption='Generated image') | |
| elif choose == 'Canvas': | |
| st.info("Step 1a. Draw your image with canvas" | |
| " \n Step 1b. You also can upload image directly by select Upload in side bar" | |
| " \n Step 2. Input prompt to instruct model (You can also change some config with advanced option if need)" | |
| " \n Step 3. Generate and enjoy") | |
| with st.form(key='canvas_generate_form'): | |
| # Specify canvas parameters in application | |
| stroke_width = st.sidebar.slider("Stroke width: ", 1, 25, 3) | |
| stroke_color = st.sidebar.color_picker("Stroke color hex: ") | |
| bg_color = st.sidebar.color_picker("Background color hex: ", "#eee") | |
| realtime_update = st.sidebar.checkbox("Update in realtime", True) | |
| # Create a canvas component | |
| col31, col32 = st.columns(2) | |
| with col31: | |
| canvas_result = st_canvas( | |
| fill_color="rgba(255, 165, 0, 0.3)", # Fixed fill color with some opacity | |
| stroke_width=stroke_width, | |
| stroke_color=stroke_color, | |
| background_color=bg_color, | |
| background_image=None, | |
| update_streamlit=realtime_update, | |
| height=512, | |
| width=512, | |
| drawing_mode="freedraw", | |
| point_display_radius=0, | |
| key="canvas", | |
| ) | |
| prompt = st.text_input(label="Prompt", placeholder='Type your instruction') | |
| with st.expander('Advanced option', expanded=False): | |
| col41, col42 = st.columns(2) | |
| with col41: | |
| image_resolution = st.slider(label="Image Resolution", min_value=256, max_value=512, value=512, step=256) | |
| strength = st.slider(label="Control Strength", min_value=0.0, max_value=2.0, value=1.0, step=0.01) | |
| guess_mode = st.checkbox(label='Guess Mode', value=False) | |
| detect_resolution = st.slider(label="HED Resolution", min_value=128, max_value=1024, value=512, step=1) | |
| ddim_steps = st.slider(label="Steps", min_value=1, max_value=100, value=20, step=1) | |
| with col42: | |
| scale = st.slider(label="Guidance Scale", min_value=0.1, max_value=30.0, value=9.0, step=0.1) | |
| seed = st.number_input(label="Seed", min_value=-1, value=-1) | |
| eta = st.number_input(label="eta (DDIM)", value=0.0) | |
| a_prompt = st.text_input(label="Added Prompt", value='best quality, extremely detailed') | |
| n_prompt = st.text_input(label="Negative Prompt", | |
| value='longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality') | |
| # Generate image from canvas | |
| generate_button = st.form_submit_button(label='Generate Image') | |
| if generate_button: | |
| if canvas_result.image_data is not None: | |
| input_image = canvas_result.image_data | |
| with st.spinner(text=f"It may take up to 1 minute under high load. Generating images..."): | |
| results = process(input_image, prompt, a_prompt, n_prompt, 1, image_resolution, detect_resolution, ddim_steps, guess_mode, strength, scale, seed, eta) | |
| H, W, C = input_image.shape | |
| output_image = cv2.resize(results[0], (W, H), interpolation=cv2.INTER_AREA) | |
| col32.image(output_image, channels='RGB', width=None, clamp=True, caption='Generated image') | |
| # Image gallery | |
| with st.expander('Image gallery', expanded=True): | |
| col01, col02, = st.columns(2) | |
| with col01: | |
| st.image('demo/example_1.jpg', caption="Sport car") | |
| st.image('demo/example_2.jpg', caption="Dog house") | |
| st.image('demo/example_3.jpg', caption="Guitar") | |
| with col02: | |
| st.image('demo/example_4.jpg', caption="Sport car") | |
| st.image('demo/example_5.jpg', caption="Dog house") | |
| st.image('demo/example_6.jpg', caption="Guitar") | |
| if __name__ == '__main__': | |
| main() | |