Diffusersで高解像度の画像を生成する

ここ1~2ヶ月、DiffusersとModalを使って画像生成の方法を模索していました。知見がたまり、ある程度自由に画像生成できるようになったと思います。せっかくなので、コーディングでつまった箇所や工夫した点をブログにまとめることにしました。この記事では、Diffusersで高解像度の画像を生成する方法について整理します。

背景

Stable Diffusion 1.5、または2系では、生成画像のサイズは標準で512x512、せいぜい1024程度まで。それ以上の解像度の画像を作ろうとすると、人物や背景が崩れてしまいます。1920pや4kの自然な画像を作るには工夫が必要となります。例えば、下記のような方法を私は試しました。

  • Upscalerを使う
    • diffusersのupscaler
    • サードパーティ性のupscaler「Real-ESRGAN」
  • ControlNet 1.1 Tileを使う

最終的にはControlNet 1.1 Tileとupscalerを組み合わせる方法に落ち着きました。

Upscalerを使う

diffusers が提供している upscaler を利用する

diffusersではいくつかのupscalerを利用することができます。

Latent Upscaler(x2)

Stable Diffusion Latent Upscalerと呼ばれる、出力画像の解像度を2倍に高める機能が提供されています。

ベース画像を生成または用意し、Latent Upscalerのパイプラインに渡して推論を実行することで、2倍の解像度の画像が得られます。

upscaler = StableDiffusionLatentUpscalePipeline.from_pretrained(model_id, torch_dtype=torch.float16)
upscaler.to("cuda")

upscaled_image = upscaler(
    prompt=prompt,
    image=base_image, " ←ここにベース画像を渡す
    num_inference_steps=20,
    guidance_scale=0,
    generator=generator,
).images[0]

Super-Resolution(x4)

Super-Resolutionは、出力画像の解像度を4倍にする機能です。使い方はLatent upscalerとほぼ同じです。

サードパーティ製のアップスケーラー「Real-ESRGAN」を利用する

Real-ESRGANは高解像度の画像を出力するためのアプリケーションです。 利用するにはいくつかの依存ライブラリとReal-ESRGAN本体のコードを実行環境に導入する必要があります。

↓必要なライブラリ。pip installで導入。

realesrgan==0.3.0
basicsr>=1.4.2
facexlib>=0.3.0
gfpgan>=1.3.8

下記をwgetやなんらかの手段でダウンロード。

https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth
https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.1/RealESRNet_x4plus.pth
https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth
https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth
https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth

Modalではdockerを使うので、Dockerfileとrequirements.txtを用意し、デプロイ時にビルドします。

UpscalerとControlNet 1.1 Tileを組み合わせる

荒い画像を元にアップスケールする場合、顔の崩れや背景の荒い箇所がそのまま拡大されてしまいます。 それを修正する方法の一つとして、img2img 機能を使うやり方があります。一度アップスケールした画像と共に、同じプロンプトと seed を diffusers に渡して再度推論を行うことで、荒い箇所を整えることが可能です。

stable-diffusion-webui では hires.fix と呼ばれる機能で実装されていますが、2023/07/15 時点、diffusers ではまだ機能として提供されていないようでした。そのため、自力で実装しました。

試行錯誤の結果、下記手順の処理でGPUの使用率を抑えつつ、納得いく仕上がりの画像を生成することができました。

  1. txt2imgでベース画像を生成
  2. ベース画像を2~4倍に拡大(Upscalerによるアップスケールではなく、単純な画像のサイズ変更)
  3. 2の画像をControlNet Tileでimg2img推論
  4. 3の画像をUpscalerで高解像度化

下記は実装したコードの一部。全体はGitHubのコードを参照のこと。

def run_inference(
    self,
    prompt: str,
    n_prompt: str,
    height: int = 512,
    width: int = 512,
    samples: int = 1,
    batch_size: int = 1,
    steps: int = 30,
    seed: int = 1,
    upscaler: str = "",
    use_face_enhancer: bool = False,
    fix_by_controlnet_tile: bool = False,
) -> list[bytes]:

    """
    1. ベース画像の生成
    """
    max_embeddings_multiples = self._count_token(p=prompt, n=n_prompt)
    generator = torch.Generator("cuda").manual_seed(seed)
    with torch.inference_mode():
        with torch.autocast("cuda"):
            generated_images = self.pipe.text2img(
                prompt * batch_size,
                negative_prompt=n_prompt * batch_size,
                height=height,
                width=width,
                num_inference_steps=steps,
                guidance_scale=7.5,
                max_embeddings_multiples=max_embeddings_multiples,
                generator=generator,
            ).images

    base_images = generated_images

    """
    2と3. ベース画像のサイズ拡大とControlNet Tileによる整形
    """
    if fix_by_controlnet_tile:
        for image in base_images:
            image = self._resize_image(image=image, scale_factor=2)
            with torch.inference_mode():
                with torch.autocast("cuda"):
                    fixed_by_controlnet = self.controlnet_pipe(
                        prompt=prompt * batch_size,
                        negative_prompt=n_prompt * batch_size,
                        num_inference_steps=steps,
                        strength=0.3,
                        guidance_scale=7.5,
                        max_embeddings_multiples=max_embeddings_multiples,
                        generator=generator,
                        image=image,
                    ).images
        generated_images.extend(fixed_by_controlnet)
        base_images = fixed_by_controlnet

    """
    4. Upscalerによるアップスケール
    """
    if upscaler != "":
        upscaled = self._upscale(
            base_images=base_images,
            half_precision=False,
            tile=700,
            upscaler=upscaler,
            use_face_enhancer=use_face_enhancer,
        )
        generated_images.extend(upscaled)

    image_output = []
    for image in generated_images:
        with io.BytesIO() as buf:
            image.save(buf, format="PNG")
            image_output.append(buf.getvalue())

    return image_output

まとめ

下図のような画像が得られるようになりました。

お題「暮れかかる湖に佇む黒髪の女性」

↓ベース画像 暮れかかる湖に佇む黒髪の女性(ベース画像)

↓高解像度化 暮れかかる湖に佇む黒髪の女性(高解像度化)