(5-4-13)基于Stable Diffusion的文生图系统（13）图像分辨率增强

文件superresolution.py功能是实现图像超分辨率处理的 Streamlit 应用。该应用允许用户上传低分辨率图像，输入提示文本以描述希望增强的内容，并通过 Stable Diffusion 模型生成高分辨率的图像。用户可以调整生成过程中的参数，如随机种子、样本数量、尺度、DDIM 步数和噪声水平等。程序通过将模型加载、图像处理和最终结果展示整合在一起，提供了一个简洁的界面来实现图像质量的提升。

1.1   加载模型

函数 initialize_model 的功能是加载和初始化给定配置文件中的模型，使用指定的检查点加载权重，并将模型移动到合适的设备（CPU 或 GPU），最后返回一个 DDIM 采样器实例。

@st.cache(allow_output_mutation=True)
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

1.2   格式转换

函数 make_batch_sd 的功能是将输入图像和提示文本转换为适合模型输入的批量格式，并将其移动到指定设备，同时支持生成多个样本。

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

1.3  生成增强后的图像

函数 make_noise_augmentation 的功能是根据输入的低分辨率图像和噪声水平，利用模型的低分辨率模型生成增强后的图像，并返回增强图像和噪声水平。

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

1.4   超分辨率生成

函数 paint 的功能是执行超分辨率生成过程，使用采样器和输入图像、提示、种子、尺度、尺寸等参数生成高分辨率图像，并在生成图像上添加水印，返回处理后的图像列表。

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
            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]}
        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}
            # uncond cond
            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
    st.text(f"upscaled image shape: {result.shape}")
    return [put_watermark(Image.fromarray(img.astype(np.uint8)), wm_encoder) for img in result]

1.5  启动 Streamlit

函数 run 的功能是启动 Streamlit 应用，提供用户界面以上传图像、输入提示和调整参数，最终调用 paint 函数生成超分辨率图像并展示结果。

def run():
    st.title("Stable Diffusion Upscaling")
    # run via streamlit run scripts/demo/depth2img.py <path-tp-config> <path-to-ckpt>
    sampler = initialize_model(sys.argv[1], sys.argv[2])

    image = st.file_uploader("Image", ["jpg", "png"])
    if image:
        image = Image.open(image)
        w, h = image.size
        st.text(f"loaded input image of size ({w}, {h})")
        width, height = map(lambda x: x - x % 64, (w, h))  # resize to integer multiple of 64
        image = image.resize((width, height))
        st.text(f"resized input image to size ({width}, {height} (w, h))")
        st.image(image)

        st.write(f"\n Tip: Add a description of the object that should be upscaled, e.g.: 'a professional photograph of a cat'")
        prompt = st.text_input("Prompt", "a high quality professional photograph")

        seed = st.number_input("Seed", min_value=0, max_value=1000000, value=0)
        num_samples = st.number_input("Number of Samples", min_value=1, max_value=64, value=1)
        scale = st.slider("Scale", min_value=0.1, max_value=30.0, value=9.0, step=0.1)
        steps = st.slider("DDIM Steps", min_value=2, max_value=250, value=50, step=1)
        eta = st.sidebar.number_input("eta (DDIM)", value=0., min_value=0., max_value=1.)

        noise_level = None
        if isinstance(sampler.model, LatentUpscaleDiffusion):
            # TODO: make this work for all models
            noise_level = st.sidebar.number_input("Noise Augmentation", min_value=0, max_value=350, value=20)
            noise_level = torch.Tensor(num_samples * [noise_level]).to(sampler.model.device).long()

        t_progress = st.progress(0)
        def t_callback(t):
            t_progress.progress(min((t + 1) / steps, 1.))

        sampler.make_schedule(steps, ddim_eta=eta, verbose=True)
        if st.button("Sample"):
            result = paint(
                sampler=sampler,
                image=image,
                prompt=prompt,
                seed=seed,
                scale=scale,
                h=height, w=width, steps=steps,
                num_samples=num_samples,
                callback=t_callback,
                noise_level=noise_level,
                eta=eta
            )
            st.write("Result")
            for image in result:
                st.image(image, output_format='PNG')


if __name__ == "__main__":
    run()

1.6   图像增强测试

通过下面的命令，使用Stable Diffusion对指定图像实现超分辨率（upscaling）提升处理。

python scripts/gradio/superresolution.py configs/stable-diffusion/x4-upscaling.yaml <path-to-checkpoint>

在Gradio或Streamlit界面中，可以上传要进行超分辨率处理的图像。如果您望处理合成的示例，建议将 noise_level 设置为较高的值（例如100）。执行后将看到超分辨率处理的结果，如图5-8所示，我们可以下载生成的高分辨率图像。

图5-8  图像超分辨率提升处理