|
import sys |
|
import torch |
|
import numpy as np |
|
import gradio as gr |
|
from PIL import Image |
|
from omegaconf import OmegaConf |
|
from einops import repeat, rearrange |
|
from pytorch_lightning import seed_everything |
|
from imwatermark import WatermarkEncoder |
|
|
|
from scripts.txt2img import put_watermark |
|
from ldm.models.diffusion.ddim import DDIMSampler |
|
from ldm.models.diffusion.ddpm import LatentUpscaleDiffusion, LatentUpscaleFinetuneDiffusion |
|
from ldm.util import exists, instantiate_from_config |
|
|
|
|
|
torch.set_grad_enabled(False) |
|
|
|
|
|
def initialize_model(config, ckpt): |
|
config = OmegaConf.load(config) |
|
model = instantiate_from_config(config.model) |
|
model.load_state_dict(torch.load(ckpt)["state_dict"], strict=False) |
|
|
|
device = torch.device( |
|
"cuda") if torch.cuda.is_available() else torch.device("cpu") |
|
model = model.to(device) |
|
sampler = DDIMSampler(model) |
|
return sampler |
|
|
|
|
|
def make_batch_sd( |
|
image, |
|
txt, |
|
device, |
|
num_samples=1, |
|
): |
|
image = np.array(image.convert("RGB")) |
|
image = torch.from_numpy(image).to(dtype=torch.float32) / 127.5 - 1.0 |
|
batch = { |
|
"lr": rearrange(image, 'h w c -> 1 c h w'), |
|
"txt": num_samples * [txt], |
|
} |
|
batch["lr"] = repeat(batch["lr"].to(device=device), |
|
"1 ... -> n ...", n=num_samples) |
|
return batch |
|
|
|
|
|
def make_noise_augmentation(model, batch, noise_level=None): |
|
x_low = batch[model.low_scale_key] |
|
x_low = x_low.to(memory_format=torch.contiguous_format).float() |
|
x_aug, noise_level = model.low_scale_model(x_low, noise_level) |
|
return x_aug, noise_level |
|
|
|
|
|
def paint(sampler, image, prompt, seed, scale, h, w, steps, num_samples=1, callback=None, eta=0., noise_level=None): |
|
device = torch.device( |
|
"cuda") if torch.cuda.is_available() else torch.device("cpu") |
|
model = sampler.model |
|
seed_everything(seed) |
|
prng = np.random.RandomState(seed) |
|
start_code = prng.randn(num_samples, model.channels, h, w) |
|
start_code = torch.from_numpy(start_code).to( |
|
device=device, dtype=torch.float32) |
|
|
|
print("Creating invisible watermark encoder (see https://github.com/ShieldMnt/invisible-watermark)...") |
|
wm = "SDV2" |
|
wm_encoder = WatermarkEncoder() |
|
wm_encoder.set_watermark('bytes', wm.encode('utf-8')) |
|
with torch.no_grad(),\ |
|
torch.autocast("cuda"): |
|
batch = make_batch_sd( |
|
image, txt=prompt, device=device, num_samples=num_samples) |
|
c = model.cond_stage_model.encode(batch["txt"]) |
|
c_cat = list() |
|
if isinstance(model, LatentUpscaleFinetuneDiffusion): |
|
for ck in model.concat_keys: |
|
cc = batch[ck] |
|
if exists(model.reshuffle_patch_size): |
|
assert isinstance(model.reshuffle_patch_size, int) |
|
cc = rearrange(cc, 'b c (p1 h) (p2 w) -> b (p1 p2 c) h w', |
|
p1=model.reshuffle_patch_size, p2=model.reshuffle_patch_size) |
|
c_cat.append(cc) |
|
c_cat = torch.cat(c_cat, dim=1) |
|
|
|
cond = {"c_concat": [c_cat], "c_crossattn": [c]} |
|
|
|
uc_cross = model.get_unconditional_conditioning(num_samples, "") |
|
uc_full = {"c_concat": [c_cat], "c_crossattn": [uc_cross]} |
|
elif isinstance(model, LatentUpscaleDiffusion): |
|
x_augment, noise_level = make_noise_augmentation( |
|
model, batch, noise_level) |
|
cond = {"c_concat": [x_augment], |
|
"c_crossattn": [c], "c_adm": noise_level} |
|
|
|
uc_cross = model.get_unconditional_conditioning(num_samples, "") |
|
uc_full = {"c_concat": [x_augment], "c_crossattn": [ |
|
uc_cross], "c_adm": noise_level} |
|
else: |
|
raise NotImplementedError() |
|
|
|
shape = [model.channels, h, w] |
|
samples, intermediates = sampler.sample( |
|
steps, |
|
num_samples, |
|
shape, |
|
cond, |
|
verbose=False, |
|
eta=eta, |
|
unconditional_guidance_scale=scale, |
|
unconditional_conditioning=uc_full, |
|
x_T=start_code, |
|
callback=callback |
|
) |
|
with torch.no_grad(): |
|
x_samples_ddim = model.decode_first_stage(samples) |
|
result = torch.clamp((x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0) |
|
result = result.cpu().numpy().transpose(0, 2, 3, 1) * 255 |
|
return [put_watermark(Image.fromarray(img.astype(np.uint8)), wm_encoder) for img in result] |
|
|
|
|
|
def pad_image(input_image): |
|
pad_w, pad_h = np.max(((2, 2), np.ceil( |
|
np.array(input_image.size) / 64).astype(int)), axis=0) * 64 - input_image.size |
|
im_padded = Image.fromarray( |
|
np.pad(np.array(input_image), ((0, pad_h), (0, pad_w), (0, 0)), mode='edge')) |
|
return im_padded |
|
|
|
|
|
def predict(input_image, prompt, steps, num_samples, scale, seed, eta, noise_level): |
|
init_image = input_image.convert("RGB") |
|
image = pad_image(init_image) |
|
width, height = image.size |
|
|
|
noise_level = torch.Tensor( |
|
num_samples * [noise_level]).to(sampler.model.device).long() |
|
sampler.make_schedule(steps, ddim_eta=eta, verbose=True) |
|
result = paint( |
|
sampler=sampler, |
|
image=image, |
|
prompt=prompt, |
|
seed=seed, |
|
scale=scale, |
|
h=height, w=width, steps=steps, |
|
num_samples=num_samples, |
|
callback=None, |
|
noise_level=noise_level |
|
) |
|
return result |
|
|
|
|
|
sampler = initialize_model(sys.argv[1], sys.argv[2]) |
|
|
|
block = gr.Blocks().queue() |
|
with block: |
|
with gr.Row(): |
|
gr.Markdown("## Stable Diffusion Upscaling") |
|
|
|
with gr.Row(): |
|
with gr.Column(): |
|
input_image = gr.Image(source='upload', type="pil") |
|
gr.Markdown( |
|
"Tip: Add a description of the object that should be upscaled, e.g.: 'a professional photograph of a cat") |
|
prompt = gr.Textbox(label="Prompt") |
|
run_button = gr.Button(label="Run") |
|
with gr.Accordion("Advanced options", open=False): |
|
num_samples = gr.Slider( |
|
label="Number of Samples", minimum=1, maximum=4, value=1, step=1) |
|
steps = gr.Slider(label="DDIM Steps", minimum=2, |
|
maximum=200, value=75, step=1) |
|
scale = gr.Slider( |
|
label="Scale", minimum=0.1, maximum=30.0, value=10, step=0.1 |
|
) |
|
seed = gr.Slider( |
|
label="Seed", |
|
minimum=0, |
|
maximum=2147483647, |
|
step=1, |
|
randomize=True, |
|
) |
|
eta = gr.Number(label="eta (DDIM)", |
|
value=0.0, min=0.0, max=1.0) |
|
noise_level = None |
|
if isinstance(sampler.model, LatentUpscaleDiffusion): |
|
|
|
noise_level = gr.Number( |
|
label="Noise Augmentation", min=0, max=350, value=20, step=1) |
|
|
|
with gr.Column(): |
|
gallery = gr.Gallery(label="Generated images", show_label=False).style( |
|
grid=[2], height="auto") |
|
|
|
run_button.click(fn=predict, inputs=[ |
|
input_image, prompt, steps, num_samples, scale, seed, eta, noise_level], outputs=[gallery]) |
|
|
|
|
|
block.launch() |
|
|
