diff --git a/README.md b/README.md index b1e0cf9..0a35801 100644 --- a/README.md +++ b/README.md @@ -136,12 +136,18 @@ To augment the well-established [img2img](https://github.com/CompVis/stable-diff Note that the original method for image modification introduces significant semantic changes w.r.t. the initial image. If that is not desired, download our [depth-conditional stable diffusion](https://huggingface.co/stabilityai/stable-diffusion-2-depth) model and the `dpt_hybrid` MiDaS [model weights](https://github.com/intel-isl/DPT/releases/download/1_0/dpt_hybrid-midas-501f0c75.pt), place the latter in a folder `midas_models` and sample via ``` -python scripts/streamlit/depth2img.py streamlit run scripts/demo/depth2img.py configs/stable-diffusion/v2-midas-inference.yaml +python scripts/gradio/depth2img.py configs/stable-diffusion/v2-midas-inference.yaml +``` + +or + +``` +streamlit run scripts/streamlit/depth2img.py configs/stable-diffusion/v2-midas-inference.yaml ``` This method can be used on the samples of the base model itself. For example, take [this sample](assets/stable-samples/depth2img/old_man.png) generated by an anonymous discord user. -Using the [streamlit](https://streamlit.io/) script `depth2img.py`, the MiDaS model first infers a monocular depth estimate given this input, +Using the [gradio](https://gradio.app) or [streamlit](https://streamlit.io/) script `depth2img.py`, the MiDaS model first infers a monocular depth estimate given this input, and the diffusion model is then conditioned on the (relative) depth output.

diff --git a/scripts/gradio/depth2img.py b/scripts/gradio/depth2img.py new file mode 100644 index 0000000..c791a4d --- /dev/null +++ b/scripts/gradio/depth2img.py @@ -0,0 +1,184 @@ +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.util import instantiate_from_config +from ldm.models.diffusion.ddim import DDIMSampler +from ldm.data.util import AddMiDaS + +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, + model_type="dpt_hybrid" +): + image = np.array(image.convert("RGB")) + image = torch.from_numpy(image).to(dtype=torch.float32) / 127.5 - 1.0 + # sample['jpg'] is tensor hwc in [-1, 1] at this point + midas_trafo = AddMiDaS(model_type=model_type) + batch = { + "jpg": image, + "txt": num_samples * [txt], + } + batch = midas_trafo(batch) + batch["jpg"] = rearrange(batch["jpg"], 'h w c -> 1 c h w') + batch["jpg"] = repeat(batch["jpg"].to(device=device), + "1 ... -> n ...", n=num_samples) + batch["midas_in"] = repeat(torch.from_numpy(batch["midas_in"][None, ...]).to( + device=device), "1 ... -> n ...", n=num_samples) + return batch + + +def paint(sampler, image, prompt, t_enc, seed, scale, num_samples=1, callback=None, + do_full_sample=False): + device = torch.device( + "cuda") if torch.cuda.is_available() else torch.device("cpu") + model = sampler.model + seed_everything(seed) + + 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) + z = model.get_first_stage_encoding(model.encode_first_stage( + batch[model.first_stage_key])) # move to latent space + c = model.cond_stage_model.encode(batch["txt"]) + c_cat = list() + for ck in model.concat_keys: + cc = batch[ck] + cc = model.depth_model(cc) + depth_min, depth_max = torch.amin(cc, dim=[1, 2, 3], keepdim=True), torch.amax(cc, dim=[1, 2, 3], + keepdim=True) + display_depth = (cc - depth_min) / (depth_max - depth_min) + depth_image = Image.fromarray( + (display_depth[0, 0, ...].cpu().numpy() * 255.).astype(np.uint8)) + cc = torch.nn.functional.interpolate( + cc, + size=z.shape[2:], + mode="bicubic", + align_corners=False, + ) + depth_min, depth_max = torch.amin(cc, dim=[1, 2, 3], keepdim=True), torch.amax(cc, dim=[1, 2, 3], + keepdim=True) + cc = 2. * (cc - depth_min) / (depth_max - depth_min) - 1. + c_cat.append(cc) + c_cat = torch.cat(c_cat, dim=1) + # cond + cond = {"c_concat": [c_cat], "c_crossattn": [c]} + + # uncond cond + uc_cross = model.get_unconditional_conditioning(num_samples, "") + uc_full = {"c_concat": [c_cat], "c_crossattn": [uc_cross]} + if not do_full_sample: + # encode (scaled latent) + z_enc = sampler.stochastic_encode( + z, torch.tensor([t_enc] * num_samples).to(model.device)) + else: + z_enc = torch.randn_like(z) + # decode it + samples = sampler.decode(z_enc, cond, t_enc, unconditional_guidance_scale=scale, + unconditional_conditioning=uc_full, callback=callback) + 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 [depth_image] + [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, strength): + init_image = input_image.convert("RGB") + image = pad_image(init_image) # resize to integer multiple of 32 + + sampler.make_schedule(steps, ddim_eta=eta, verbose=True) + assert 0. <= strength <= 1., 'can only work with strength in [0.0, 1.0]' + do_full_sample = strength == 1. + t_enc = min(int(strength * steps), steps-1) + result = paint( + sampler=sampler, + image=image, + prompt=prompt, + t_enc=t_enc, + seed=seed, + scale=scale, + num_samples=num_samples, + callback=None, + do_full_sample=do_full_sample + ) + return result + + +sampler = initialize_model(sys.argv[1], sys.argv[2]) + +block = gr.Blocks().queue() +with block: + with gr.Row(): + gr.Markdown("## Stable Diffusion Depth2Img") + + with gr.Row(): + with gr.Column(): + input_image = gr.Image(source='upload', type="pil") + prompt = gr.Textbox(label="Prompt") + run_button = gr.Button(label="Run") + with gr.Accordion("Advanced options", open=False): + num_samples = gr.Slider( + label="Images", minimum=1, maximum=4, value=1, step=1) + ddim_steps = gr.Slider(label="Steps", minimum=1, + maximum=50, value=50, step=1) + scale = gr.Slider( + label="Guidance Scale", minimum=0.1, maximum=30.0, value=9.0, step=0.1 + ) + strength = gr.Slider( + label="Strength", minimum=0.0, maximum=1.0, value=0.9, step=0.01 + ) + seed = gr.Slider( + label="Seed", + minimum=0, + maximum=2147483647, + step=1, + randomize=True, + ) + eta = gr.Number(label="eta (DDIM)", value=0.0) + 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, ddim_steps, num_samples, scale, seed, eta, strength], outputs=[gallery]) + + +block.launch()