app.py

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: <a href=\"https://github.com/robin025\">Robin Singh</a>  """)

demo.queue()
demo.launch(debug=True, share=True, height=768)