from diffusers import StableDiffusionPipeline, StableDiffusionImg2ImgPipeline, StableDiffusionUpscalePipeline, DiffusionPipeline, DPMSolverMultistepScheduler,LMSDiscreteScheduler,DDIMScheduler,EulerDiscreteScheduler,PNDMScheduler,DDPMScheduler,EulerAncestralDiscreteScheduler import gradio as gr import torch from PIL import Image import random state = None current_steps = 25 # SD 2.1 is used model_id = 'stabilityai/stable-diffusion-2-1' # Schedulers Used DPMS = DPMSolverMultistepScheduler.from_pretrained(model_id, subfolder="scheduler") EADS = EulerAncestralDiscreteScheduler.from_pretrained(model_id, subfolder="scheduler") LMSD = LMSDiscreteScheduler.from_pretrained(model_id, subfolder="scheduler") DDIM = DDIMScheduler.from_pretrained(model_id, subfolder="scheduler") EDS = EulerDiscreteScheduler.from_pretrained(model_id, subfolder="scheduler") PNMS = PNDMScheduler.from_pretrained(model_id, subfolder="scheduler") DDPM = DDPMScheduler.from_pretrained(model_id, subfolder="scheduler") scheduler_types={ "DPMS":DPMSolverMultistepScheduler.from_pretrained(model_id, subfolder="scheduler"), "EADS":EulerAncestralDiscreteScheduler.from_pretrained(model_id, subfolder="scheduler"), "LMSD":LMSDiscreteScheduler.from_pretrained(model_id, subfolder="scheduler"), "DDIM":DDIMScheduler.from_pretrained(model_id, subfolder="scheduler"), "EDS":EulerDiscreteScheduler.from_pretrained(model_id, subfolder="scheduler"), "PNMS":PNDMScheduler.from_pretrained(model_id, subfolder="scheduler"), "DDPM":DDPMScheduler.from_pretrained(model_id, subfolder="scheduler"), } # Creating Simple Customized pipeline pipe = StableDiffusionPipeline.from_pretrained( model_id, revision="fp16", torch_dtype=torch.float16, scheduler=DPMS ).to("cuda") pipe.enable_attention_slicing() # pipe.enable_xformers_memory_efficient_attention() # Different Pipeline states pipe_i2i = None pipe_upscale = None pipe_inpaint = None attn_slicing_enabled = True mem_eff_attn_enabled = False # Different Modes of Inference (VideoGen : TODO) modes = { 'txt2img': 'Text to Image', 'img2img': 'Image to Image', 'inpaint': 'Inpainting', 'upscale4x': 'Upscale', 'VideoGen':"Generation of Video" } ############################################################################### current_mode = modes['txt2img'] def error_str(error, title="Error"): return f"""#### {title} {error}""" if error else "" def update_state(new_state): global state state = new_state def update_state_info(old_state): if state and state != old_state: return gr.update(value=state) def set_mem_optimizations(pipe): if attn_slicing_enabled: pipe.enable_attention_slicing() else: pipe.disable_attention_slicing() ############################################################################### # Function for creating a new pipleline for Image to Image Generation. def get_i2i_pipe(scheduler): update_state("Loading image to image model...") pipe = StableDiffusionImg2ImgPipeline.from_pretrained( model_id, revision="fp16" if torch.cuda.is_available() else "fp32", torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32, scheduler=scheduler ) set_mem_optimizations(pipe) pipe.to("cuda") return pipe ############################################################################### # Function for creating a new pipleline for Inpaint Pipeline. def get_inpaint_pipe(): update_state("Loading inpainting model...") pipe = DiffusionPipeline.from_pretrained( "stabilityai/stable-diffusion-2-inpainting", revision="fp16" if torch.cuda.is_available() else "fp32", torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32, ).to("cuda") pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config) pipe.enable_attention_slicing() # pipe.enable_xformers_memory_efficient_attention() return pipe ############################################################################### # Function for creating a new pipleline for Upscaling the image. def get_upscale_pipe(scheduler): update_state("Loading upscale model...") pipe = StableDiffusionUpscalePipeline.from_pretrained( "stabilityai/stable-diffusion-x4-upscaler", revision="fp16" if torch.cuda.is_available() else "fp32", torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32, ) pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config) set_mem_optimizations(pipe) pipe.to("cuda") return pipe ############################################################################### def switch_attention_slicing(attn_slicing): global attn_slicing_enabled attn_slicing_enabled = attn_slicing def switch_mem_eff_attn(mem_eff_attn): global mem_eff_attn_enabled mem_eff_attn_enabled = mem_eff_attn def pipe_callback(step: int, timestep: int, latents: torch.FloatTensor): update_state(f"{step}/{current_steps} steps")#\nTime left, sec: {timestep/100:.0f}") ############################################################################### # Main Inference Function def inference(inf_mode, prompt, n_images, guidance, steps, width=768, height=768, seed=0, img=None, strength=0.5, neg_prompt="", scheduler_mode=None): update_state(" ") SDD = scheduler_types[scheduler_mode] SDD = scheduler_types.get(scheduler_mode) print(SDD) pipe.scheduler = SDD global current_mode if inf_mode != current_mode: pipe.to("cuda" if inf_mode == modes['txt2img'] else "cpu") if pipe_i2i is not None: pipe_i2i.to("cuda" if inf_mode == modes['img2img'] else "cpu") if pipe_inpaint is not None: pipe_inpaint.to("cuda" if inf_mode == modes['inpaint'] else "cpu") if pipe_upscale is not None: pipe_upscale.to("cuda" if inf_mode == modes['upscale4x'] else "cpu") current_mode = inf_mode if seed == 0: seed = random.randint(0, 2147483647) generator = torch.Generator('cuda').manual_seed(seed) prompt = prompt try: if inf_mode == modes['txt2img']: return txt_to_img(prompt, n_images, neg_prompt, guidance, steps, width, height, generator, seed), gr.update(visible=False, value=None) elif inf_mode == modes['img2img']: if img is None: return None, gr.update(visible=True, value=error_str("Image is required for Image to Image mode")) return img_to_img(prompt, n_images, neg_prompt, img, strength, guidance, steps, width, height, generator, seed), gr.update(visible=False, value=None) elif inf_mode == modes['inpaint']: if img is None: return None, gr.update(visible=True, value=error_str("Image is required for Inpainting mode")) return inpaint(prompt, n_images, neg_prompt, img, guidance, steps, width, height, generator, seed), gr.update(visible=False, value=None) elif inf_mode == modes['upscale4x']: if img is None: return None, gr.update(visible=True, value=error_str("Image is required for Upscale mode")) return upscale(prompt, n_images, neg_prompt, img, guidance, steps, generator), gr.update(visible=False, value=None) # elif inf_mode == modes['VideoGen']: # if img is None: # return None, gr.update(visible=True, value=error_str("Image is required for Video Generation")) # return upscale(prompt, n_images, neg_prompt, img, guidance, steps, generator, seed), gr.update(visible=False, value=None) except Exception as e: return None, gr.update(visible=True, value=error_str(e)) ############################################################################### # Text To Image def txt_to_img(prompt, n_images, neg_prompt, guidance, steps, width, height, generator, seed): result = pipe( prompt, num_images_per_prompt = n_images, negative_prompt = neg_prompt, num_inference_steps = int(steps), guidance_scale = guidance, width = width, height = height, generator = generator, callback=pipe_callback).images update_state(f"Done. Seed: {seed}") return result ############################################################################### # Image To image def img_to_img(prompt, n_images, neg_prompt, img, strength, guidance, steps, width, height, generator, seed): global pipe_i2i if pipe_i2i is None: pipe_i2i = get_i2i_pipe(DPMS) img = img['image'] ratio = min(height / img.height, width / img.width) img = img.resize((int(img.width * ratio), int(img.height * ratio)), Image.LANCZOS) result = pipe_i2i( prompt, num_images_per_prompt = n_images, negative_prompt = neg_prompt, image = img, num_inference_steps = int(steps), strength = strength, guidance_scale = guidance, # width = width, # height = height, generator = generator, callback=pipe_callback).images update_state(f"Done. Seed: {seed}") return result ############################################################################### # Inpaint def inpaint(prompt, n_images, neg_prompt, img, guidance, steps, width, height, generator, seed): global pipe_inpaint if pipe_inpaint is None: pipe_inpaint = get_inpaint_pipe() inp_img = img['image'] mask = img['mask'] inp_img = square_padding(inp_img) mask = square_padding(mask) inp_img = inp_img.resize((512, 512)) mask = mask.resize((512, 512)) result = pipe_inpaint( prompt, image = inp_img, mask_image = mask, num_images_per_prompt = n_images, negative_prompt = neg_prompt, num_inference_steps = int(steps), guidance_scale = guidance, generator = generator, callback=pipe_callback).images update_state(f"Done. Seed: {seed}") return result def square_padding(img): width, height = img.size if width == height: return img new_size = max(width, height) new_img = Image.new('RGB', (new_size, new_size), (0, 0, 0, 255)) new_img.paste(img, ((new_size - width) // 2, (new_size - height) // 2)) return new_img ############################################################################### # Upscale def upscale(prompt, n_images, neg_prompt, img, guidance, steps, generator): global pipe_upscale if pipe_upscale is None: pipe_upscale = get_upscale_pipe(DPMS) img = img['image'] return upscale_tiling(prompt, neg_prompt, img, guidance, steps, generator) ############################################################################### # Upscale def upscale_tiling(prompt, neg_prompt, img, guidance, steps, generator): width, height = img.size # calculate the padding needed to make the image dimensions a multiple of 128 padding_x = 128 - (width % 128) if width % 128 != 0 else 0 padding_y = 128 - (height % 128) if height % 128 != 0 else 0 # create a white image of the right size to be used as padding padding_img = Image.new('RGB', (padding_x, padding_y), color=(255, 255, 255, 0)) # paste the padding image onto the original image to add the padding img.paste(padding_img, (width, height)) # update the image dimensions to include the padding width += padding_x height += padding_y if width > 128 or height > 128: num_tiles_x = int(width / 128) num_tiles_y = int(height / 128) upscaled_img = Image.new('RGB', (img.size[0] * 4, img.size[1] * 4)) for x in range(num_tiles_x): for y in range(num_tiles_y): update_state(f"Upscaling tile {x * num_tiles_y + y + 1}/{num_tiles_x * num_tiles_y}") tile = img.crop((x * 128, y * 128, (x + 1) * 128, (y + 1) * 128)) upscaled_tile = pipe_upscale( prompt="", image=tile, num_inference_steps=steps, guidance_scale=guidance, generator=generator, ).images[0] upscaled_img.paste(upscaled_tile, (x * upscaled_tile.size[0], y * upscaled_tile.size[1])) return [upscaled_img] else: return pipe_upscale( prompt=prompt, image=img, num_inference_steps=steps, guidance_scale=guidance, negative_prompt = neg_prompt, generator=generator, ).images # Mode Change def on_mode_change(mode): return gr.update(visible = mode in (modes['img2img'], modes['inpaint'], modes['upscale4x'])), \ gr.update(visible = mode == modes['inpaint']), \ gr.update(visible = mode == modes['upscale4x']), \ gr.update(visible = mode == modes['img2img']) def on_steps_change(steps): global current_steps current_steps = steps ############################################################################### # Gradio UI css = """#primary {color: yellow} #main-div {color:#2B0230} .main-div div{display:flex;flex-direction:column;align-items:center;gap:.8rem;font-size:1.75rem}.main-div div h1{font-weight:900;margin-bottom:7px}.main-div p{margin-bottom:10px;font-size:94%}a{text-decoration:underline}.tabs{margin-top:0;margin-bottom:0}#gallery{min-height:20rem} """ with gr.Blocks(css=css) as demo: gr.HTML( f""" Genie : Stable Diffusion """ ) with gr.Row(elem_id='main-div'): with gr.Column(scale=100): inf_mode = gr.Radio(label="Modes", choices=list(modes.values())[:4], value=modes['txt2img']) # TODO remove [:3] limit with gr.Group(visible=False) as i2i_options: image = gr.Image(label="Image", height=128, type="pil", tool='sketch') inpaint_info = gr.Markdown("Inpainting resizes and pads images to 512x512", visible=False) upscale_info = gr.Markdown("""Best for small images (128x128 or smaller). Bigger images will be sliced into 128x128 tiles which will be upscaled individually. This is done to avoid running out of GPU memory.""", visible=False) videogen_info = gr.Markdown(""" Video Generation : TODO """) strength = gr.Slider(label="Transformation strength", minimum=0, maximum=1, step=0.01, value=0.5) with gr.Group(): neg_prompt = gr.Textbox(label="Negative prompt", placeholder="What to exclude from the image") choose_scheduler = gr.Dropdown(["DPMS","EADS","LMSD","DDIM","EDS","PNMS","DDPM"]) n_images = gr.Slider(label="Number of images", value=1, minimum=1, maximum=10, step=1) with gr.Row(): guidance = gr.Slider(label="Guidance scale", value=7.5, maximum=15) steps = gr.Slider(label="Steps", value=current_steps, minimum=1, maximum=100, step=.5) with gr.Row(): width = gr.Slider(label="Width", value=768, minimum=64, maximum=1024, step=8) height = gr.Slider(label="Height", value=768, minimum=64, maximum=1024, step=8) with gr.Column(scale=100): with gr.Group(): with gr.Row(): prompt = gr.Textbox(label="Prompt", show_label=False, max_lines=2,placeholder=f"enter something").style(container=True) gallery = gr.Gallery(label="Generated images", show_label=False).style(grid=[2], height="auto") state_info = gr.Textbox(label="State", show_label=False, max_lines=2).style(container=False) generate = gr.Button(value="Generate", elem_id="primary").style(rounded=(False, True, True, False),) error_output = gr.Markdown(visible=False) with gr.Row(): with gr.Column(scale=100): seed = gr.Slider(0, 2147483647, label='Seed', value=456785, step=1) with gr.Accordion("Memory optimization"): attn_slicing = gr.Checkbox(label="Attention slicing", value=attn_slicing_enabled) inf_mode.change(on_mode_change, inputs=[inf_mode], outputs=[i2i_options, inpaint_info, upscale_info, strength], queue=False) steps.change(on_steps_change, inputs=[steps], outputs=[], queue=False) attn_slicing.change(lambda x: switch_attention_slicing(x), inputs=[attn_slicing], queue=False) inputs = [inf_mode, prompt, n_images, guidance, steps, width, height, seed, image, strength, neg_prompt,choose_scheduler] outputs = [gallery, error_output] prompt.submit(inference, inputs=inputs, outputs=outputs) generate.click(inference, inputs=inputs, outputs=outputs) demo.load(update_state_info, inputs=state_info, outputs=state_info, every=0.5, show_progress=False) gr.HTML(""" Developed by: Robin Singh """) demo.queue() demo.launch(debug=True, share=True, height=768)