stable-diffusion-finetune/scripts/demo/inpainting.py

209 lines
7.4 KiB
Python
Raw Normal View History

2022-07-26 10:24:02 +02:00
import streamlit as st
import torch
import cv2
import numpy as np
from ldm.util import instantiate_from_config
from omegaconf import OmegaConf
from PIL import Image
import io
from streamlit_drawable_canvas import st_canvas
torch.set_grad_enabled(False)
def sample(
model,
prompt,
n_runs=3,
n_samples=2,
H=512,
W=512,
scale=10.0,
ddim_steps=50,
callback=None,
image=None,
mask=None,
):
batch = np2batch(image=image, mask=mask, txt=prompt)
self = model
unconditional_guidance_scale = scale
unconditional_guidance_label = [""]
use_ddim = True
ddim_eta = 0
N = 1
ema_scope = self.ema_scope
z, c, x, xrec, xc = self.get_input(batch, self.first_stage_key, bs=N, return_first_stage_outputs=True)
c_cat, c = c["c_concat"][0], c["c_crossattn"][0]
if unconditional_guidance_scale > 1.0:
uc_cross = self.get_unconditional_conditioning(N, unconditional_guidance_label)
uc_cat = c_cat
uc_full = {"c_concat": [uc_cat], "c_crossattn": [uc_cross]}
with ema_scope("Sampling with classifier-free guidance"):
samples_cfg, _ = self.sample_log(cond={"c_concat": [c_cat], "c_crossattn": [c]},
batch_size=N, ddim=use_ddim,
ddim_steps=ddim_steps, eta=ddim_eta,
unconditional_guidance_scale=unconditional_guidance_scale,
unconditional_conditioning=uc_full,
)
samples = self.decode_first_stage(samples_cfg)
else:
raise ValueError()
samples = torch2np(samples)
return samples
def np2batch(
image,
mask,
txt):
print("###")
print(image.shape)
print(mask.shape)
print("###")
# image hwc in -1 1
image = torch.from_numpy(image).to(dtype=torch.float32)/127.5-1.0
mask[mask < 0.5] = 0
mask[mask > 0.5] = 1
mask = torch.from_numpy(mask)[:,:,:1]
masked_image = image * (mask < 0.5)
batch = {
"jpg": image[None],
"txt": [txt],
"mask": mask[None],
"masked_image": masked_image[None],
}
return batch
def torch2np(x):
x = ((x+1.0)*127.5).clamp(0, 255).to(dtype=torch.uint8)
x = x.permute(0, 2, 3, 1).detach().cpu().numpy()
return x
@st.cache(allow_output_mutation=True)
def init():
state = dict()
return state
def load_model_from_config(config, ckpt, verbose=False):
print(f"Loading model from {ckpt}")
pl_sd = torch.load(ckpt, map_location="cpu")
global_step = pl_sd.get("global_step", "?")
sd = pl_sd["state_dict"]
model = instantiate_from_config(config.model)
m, u = model.load_state_dict(sd, strict=False)
if len(m) > 0 and verbose:
print("missing keys:")
print(m)
if len(u) > 0 and verbose:
print("unexpected keys:")
print(u)
model.cuda()
model.eval()
print(f"Loaded global step {global_step}")
return model
if __name__ == "__main__":
st.title("Stable Inpainting")
state = init()
if not "model" in state:
config = "/fsx/stable-diffusion/stable-diffusion/configs/stable-diffusion/inpainting/v1-finetune-for-inpainting-laion-iaesthe.yaml"
ckpt = "/fsx/robin/stable-diffusion/stable-diffusion/logs/2022-07-24T16-01-27_v1-finetune-for-inpainting-laion-iaesthe/checkpoints/last.ckpt"
config = OmegaConf.load(config)
model = load_model_from_config(config, ckpt)
state["model"] = model
uploaded_file = st.file_uploader("Upload image to inpaint")
if uploaded_file is not None:
image = Image.open(io.BytesIO(uploaded_file.getvalue()))
width, height = image.size
smaller = min(width, height)
crop = (
(width-smaller)//2,
(height-smaller)//2,
(width-smaller)//2+smaller,
(height-smaller)//2+smaller,
)
image = image.crop(crop)
image = image.resize((512, 512))
#st.write("Uploaded Image")
#st.image(image)
st.write("Draw a mask (and send it to streamlit, button lower left)")
stroke_width = int(st.number_input("Stroke Width", value=50))
canvas_result = st_canvas(
fill_color="rgba(255, 255, 255)", # Fixed fill color with some opacity
stroke_width=stroke_width,
stroke_color="rgb(0, 0, 0)",
background_color="rgb(0, 0, 0)",
background_image=image if image is not None else Image.fromarray(255*np.ones((512,512,3),
dtype=np.uint8)),
update_streamlit=False,
height=image.size[1] if image is not None else 512,
width=image.size[0] if image is not None else 512,
drawing_mode="freedraw",
point_display_radius=0,
key="canvas",
)
if canvas_result:
mask = canvas_result.image_data
mask = np.array(mask)[:,:,[3,3,3]]
mask = mask > 127
# visualize
bdry = cv2.dilate(mask.astype(np.uint8), np.ones((3,3), dtype=np.uint8))
bdry = (bdry > 0) & ~mask
masked_image = np.array(image)*(1-mask) + mask*0.3*np.array(image)
masked_image[:,:,0][bdry[:,:,0]] = 255
masked_image[:,:,1][bdry[:,:,1]] = 0
masked_image[:,:,2][bdry[:,:,2]] = 0
st.write("Masked Image")
st.image(Image.fromarray(masked_image.astype(np.uint8)))
prompt = st.text_input("Prompt")
scale = float(st.number_input("Guidance", value=10.0))
t_total = int(st.number_input("Diffusion steps", value=50))
if st.button("Sample"):
st.text("Sampling")
batch_progress = st.progress(0)
batch_total = 3
t_progress = st.progress(0)
result = st.empty()
#canvas = make_canvas(2, 3)
def callback(x, batch, t):
#result.text(f"{batch}, {t}")
batch_progress.progress(min(1.0, (batch+1)/batch_total))
t_progress.progress(min(1.0, (t+1)/t_total))
update_canvas(canvas, x, batch)
result.image(canvas)
samples = sample(
state["model"],
prompt,
n_runs=3,
n_samples=2,
H=512,
W=512,
scale=scale,
ddim_steps=t_total,
callback=callback,
image=np.array(image),
mask=np.array(mask),
)
st.text("Samples")
st.image(samples[0])