File size: 4,230 Bytes
c25e2cc 3cdacdf 6255790 5e25b83 6255790 c25e2cc 6255790 e79152d 6255790 e79152d 6255790 5e25b83 9b96547 f55706c b592002 6255790 6505e1f e200bfb 6255790 9b96547 6255790 5e25b83 e200bfb 5e25b83 1a248f3 aa5232a 5e25b83 6255790 e200bfb c0aaf41 e200bfb 5e25b83 6255790 52d8a48 5e25b83 6255790 5e25b83 6255790 5e25b83 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 |
import gradio as gr
import torch
import requests
from io import BytesIO
from diffusers import StableDiffusionPipeline
from diffusers import DDIMScheduler
from utils import *
from inversion_utils import *
from modified_pipeline_semantic_stable_diffusion import SemanticStableDiffusionPipeline
from torch import autocast, inference_mode
def invert(x0, prompt_src="", num_diffusion_steps=100, cfg_scale_src = 3.5, eta = 1):
# inverts a real image according to Algorihm 1 in https://arxiv.org/pdf/2304.06140.pdf,
# based on the code in https://github.com/inbarhub/DDPM_inversion
# returns wt, zs, wts:
# wt - inverted latent
# wts - intermediate inverted latents
# zs - noise maps
sd_pipe.scheduler.set_timesteps(num_diffusion_steps)
# vae encode image
with autocast("cuda"), inference_mode():
w0 = (sd_pipe.vae.encode(x0).latent_dist.mode() * 0.18215).float()
# find Zs and wts - forward process
wt, zs, wts = inversion_forward_process(sd_pipe, w0, etas=eta, prompt=prompt_src, cfg_scale=cfg_scale_src, prog_bar=True, num_inference_steps=num_diffusion_steps)
return wt, zs, wts
def sample(wt, zs, wts, prompt_tar="", cfg_scale_tar=15, skip=36, eta = 1):
# reverse process (via Zs and wT)
w0, _ = inversion_reverse_process(sd_pipe, xT=wts[skip], etas=eta, prompts=[prompt_tar], cfg_scales=[cfg_scale_tar], prog_bar=True, zs=zs[skip:])
# vae decode image
with autocast("cuda"), inference_mode():
x0_dec = sd_pipe.vae.decode(1 / 0.18215 * w0).sample
if x0_dec.dim()<4:
x0_dec = x0_dec[None,:,:,:]
img = image_grid(x0_dec)
return img
# load pipelines
sd_model_id = "runwayml/stable-diffusion-v1-5"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
sd_pipe = StableDiffusionPipeline.from_pretrained(sd_model_id).to(device)
sd_pipe.scheduler = DDIMScheduler.from_config(sd_model_id, subfolder = "scheduler")
sem_pipe = SemanticStableDiffusionPipeline.from_pretrained(sd_model_id, torch_dtype=torch.float16).to(device)
def edit(input_image, input_image_prompt='', target_prompt='', edit_prompt='',
negative_guidance = False, edit_warmup_steps=5,
edit_guidance_scale=8, guidance_scale=15, skip=36, num_diffusion_steps=100,
):
offsets=(0,0,0,0)
x0 = load_512(input_image, *offsets, device)
# invert
wt, zs, wts = invert(x0 =x0 , prompt_src=input_image_prompt, num_diffusion_steps=num_diffusion_steps)
latnets = wts[skip].expand(1, -1, -1, -1)
eta = 1
#pure DDPM output
pure_ddpm_out = sample(wt, zs, wts, prompt_tar=target_prompt,
cfg_scale_tar=guidance_scale, skip=skip,
eta = eta)
editing_args = dict(
editing_prompt = [edit_prompt],
reverse_editing_direction = [negative_guidance],
edit_warmup_steps=[edit_warmup_steps],
edit_guidance_scale=[edit_guidance_scale],
edit_threshold=[.93],
edit_momentum_scale=0.5,
edit_mom_beta=0.6
)
sega_out = sem_pipe(prompt=target_prompt,eta=eta, latents=latnets,
num_images_per_prompt=1,
num_inference_steps=num_diffusion_steps,
use_ddpm=True, wts=wts, zs=zs[skip:], **editing_args)
return pure_ddpm_out,sega_out.images[0]
# See the gradio docs for the types of inputs and outputs available
inputs = [
gr.Image(label="input image", shape=(512, 512)),
gr.Textbox(label="input prompt"),
gr.Textbox(label="target prompt"),
gr.Textbox(label="SEGA edit concept"),
gr.Checkbox(label="SEGA negative_guidance"),
gr.Slider(label="warmup steps", minimum=1, maximum=30, value=5),
gr.Slider(label="edit guidance scale", minimum=0, maximum=15, value=3.5),
gr.Slider(label="guidance scale", minimum=7, maximum=18, value=15),
gr.Slider(label="skip", minimum=0, maximum=40, value=36),
gr.Slider(label="num diffusion steps", minimum=0, maximum=300, value=100)
]
outputs = [gr.Image(label="DDPM"),gr.Image(label="DDPM+SEGA")]
# And the minimal interface
demo = gr.Interface(
fn=edit,
inputs=inputs,
outputs=outputs,
)
demo.launch() # debug=True allows you to see errors and output in Colab |