DiffBIRで劣化した画像の復元を行う手順（Windows）

2023年9月8日 04:11

Last update 09-08-2023
※再チェックを行っていないので、誤り等があればお知らせください。

▼ 本記事について

概要

　DiffBIRは、中国科学院深セン先進技術研究院（Shenzhen Institute of Advanced Technology, Chinese Academy of Sciences）のXinqi Lin氏、上海人工智能实验室（Shanghai AI Laboratory）のJingwen He氏らにより提案された画像復元の手法で、下記の特徴があります。

特定の劣化プロセスに限定しない復元が可能。
Stage1のノイズ除去と、Stage 2のStable Diffusionによる補完のセット。解像度の変更も可能。
従来の画像超解像（BSR）＋ブラインド顔復元（BFR）よりも優れた性能を示す。
復元に使用するモデルは、general image用とface用の二種類がある。
SwinIR以外の超解像の利用や、モデルのトレーニングが可能（リポジトリの説明を参照）。

　本記事では、DiffBIRをインストールして利用するための手順について説明します。

動作環境

　VRAM 12GB以上のGeForce RTX（20シリーズ以降）、Python 3.10、CUDA Toolkit 11.8が導入済みであることを前提とします。起動のみで7GB少々のVRAMを消費するため、8GBでは著しい速度の低下が予想されます。

公式リポジトリ、論文等

DiffBIR: Towards Blind Image Restoration with Generative Diffusion Prior
https://github.com/XPixelGroup/DiffBIR
https://arxiv.org/abs/2308.15070

Now we have DiffBIR Colab 🥳 Thanks to Xinqi Lin ❤ Jingwen He ❤ Ziyan Chen ❤ Zhaoyang Lyu ❤ Ben Fei ❤ @doubledaibo ❤ Wanli Ouyang ❤ Yu Qiao ❤ Chao Dong ❤

📄paper: https://t.co/qeiNIsP0mT
🧬code: https://t.co/SF0M2LyJIc
🦒colab: please try it 🐣 https://t.co/7lAVOnoIMa pic.twitter.com/w1t7Bgjjcl
— camenduru (@camenduru) September 6, 2023

▼ インストール

　リポジトリの説明はAnacondaを前提としているため、Pythonでインストールできるようアレンジを行いました。手順はあまり洗練されていませんがご容赦ください。

ダウンロードと仮想環境の作成

　作業ディレクトリを「C:\aiwork」としていますので、適宜読み替えてください。コマンドプロンプトを開いて、ディレクトリ「C:\aiwork\DiffBIR」が無い状態で下記のコマンドを順に実行してください。

　Gitを使用しないでダウンロードしたい方のために、「git clone ～」のコマンドの代わりの操作を説明します。Webブラウザで公式リポジトリへアクセスして、右側の「Code」→「Download」の順にクリックしてください。その後、ダウンロードしたファイルを適当なディレクトリに解凍してください。

cd \aiwork

git clone https://github.com/XPixelGroup/DiffBIR
cd DiffBIR

python -m venv venv
venv\Scripts\activate
（ここで行頭に (venv) が付いていることを確認する）

python -m pip install --upgrade pip

パッケージのインストール1

　下記のコマンドを順に実行してください。特に1行目は、パッケージのダウンロード等でしばらくかかります。

pip install torch torchvision --extra-index-url https://download.pytorch.org/whl/cu118
pip install xformers

　公式リポジトリの説明からは「torch==1.12.1」「xformers==0.0.16」等のバージョン指定が読み取れますが、特に指定しなくても動作したのでこのままにしています。

　CUDA 11.7以下向けのPyTorchを使用したい場合は、上記の「cu118」の箇所を「cu117」「cu116」「cu113」のいずれかに書き換えてください。参考まで、デフォルトでインストールされるバージョンを記載します。

cu113の場合
torch==1.12.1 torchvision==0.13.1 

cu116の場合
torch==1.13.1 torchvision==0.14.1
（torch==1.12.1 torchvision==0.13.1 も指定可）

cu117の場合
torch==1.13.1 torchvision==0.14.1

cu118の場合（現時点）
torch==2.0.1 torchvision==0.15.2

パッケージのインストール2

　続いて install_env.sh の内容を参考に、必要なパッケージをインストールします。新規ファイル「C:\aiwork\DiffBIR\requirements.txt」を下記の内容で作成してください。

pytorch_lightning==1.4.2
einops
transformers
chardet
open-clip-torch
omegaconf
torchmetrics==0.6.0
opencv-python-headless
scipy
matplotlib
lpips
gradio

　続けて、下記のコマンドを順に実行してください。Tritonは標準のパッケージが用意されていないため、ビルド済みのものを拝借します（→参照）。

pip install -r requirements.txt
pip install https://huggingface.co/r4ziel/xformers_pre_built/resolve/main/triton-2.0.0-cp310-cp310-win_amd64.whl

モデルのダウンロード

　使用するモデルは https://huggingface.co/lxq007/DiffBIR/tree/main にあります。今後の手順のため、最低でも「general_full_v1.ckpt」と「general_swinir_v1.ckpt」の二つをダウンロードして、「C:\aiwork\DiffBIR」に移動してください。

▼ 実行（動作確認）

起動

　リポジトリの「Quick Start」にならって動作を確認してみます。先ほどの手順から続けて、下記のコマンドを実行してください。初回は少しだけ待たされる可能性があります。

python gradio_diffbir.py --ckpt ./general_full_v1.ckpt --config configs/model/cldm.yaml --reload_swinir --swinir_ckpt ./general_swinir_v1.ckpt

　起動が完了すると、下記の表示で停止しますので、Webブラウザで http://127.0.0.1:7860/ へアクセスしてください。

Running on local URL:  http://0.0.0.0:7860

To create a public link, set `share=True` in `launch()`.

復元の実行

　512x512の画像（230908a_512.jpg）をダウンロードしてください。劣化させるため、低解像度で圧縮率を上げて保存し、それを読み込んで解像度を上げています。他の画像でも構いませんが、同じような手順で劣化したイメージを用意してください。解像度は256x256または512x512が良いでしょう。

　次に、UIの左側の欄へ画像をドラッグ＆ドロップしてから「Run」をクリックしてください。256x256の画像の場合は、「SR Scale」を2に設定すると512x512で出力されます。

　処理が完了すると右側に表示されます。

　上記の画像のように、右上にダウンロードのボタンが無い場合は、右側の画像を一度クリックしてください。

　GUIで実行した場合、画像を自動保存する機能は無いように見えます。リポジトリの「Inference」にて、コマンドライン上で「inference.py」を実行して、任意のディレクトリに保存する方法が記載されています。

終了する

　コマンドプロンプトのウインドウで「Ctrl + C」を押すか、おもむろにコマンドプロンプトの画面を閉じてください。作業のためコマンドプロンプトの画面を残したい場合は必要に応じて、下記のコマンドを実行して仮想環境から抜けてください。行頭の「(venv)」の表示が消えます。

deactivate

▼ 解像度等に関する説明（重要）

　結論を先に挙げると、「Keep Original Size」はON（デフォルト）のままで、ディティールの細かさを「Image size」で、出力の倍率を「SR Scale」で設定する方法が簡単です。

　それでは、入力解像度、出力解像度、「SR Scale」、「Image size」、「Keep Original Size」の関係について説明します。なお、正方形の画像のみで確認を行いました。

Keep Original SizeがONの場合

　出力解像度は、入力解像度の「SR Scale」倍です。ディティールの細かさと出力にかかる時間は、「Image size」に依存します。

Keep Original SizeがOFFの場合

　出力解像度は、入力解像度の「SR Scale」倍と「Image size」の、大きい方の値です。ディティールの細かさと出力にかかる時間は、出力解像度に依存します。

考察

　設定を変えて出力を見比べた結果、復元実行時の挙動は下記のようになっていると考えられます。

Keep Original SizeがOFFの場合は、内部で生成したイメージをそのまま出力する。解像度は「入力画像のSR Scale倍」と「Image size」の、大きい方の値に決定する。
Keep Original SizeがONの場合は、OFFの場合の動作に加えて、内部で生成したイメージを入力解像度の「SR Scale」倍にリサイズする。

　根拠として、下記の点を挙げます。「Keep=Keep Original Size、入力=入力解像度(px)、SR=SR Scale、size=Image size(px)、出力=出力解像度(px)」とします。

Keep=ON、入力=128、SR=2のとき、出力=256固定だが、ディティールの細かさと出力にかかる時間がsize（256, 512, 768にて確認）に依存している。
Keep=ON、入力=128、SR=1、size=768では時間がかかるが、出力=128と小さいながらも綺麗な出力が得られる。
Keep=OFF、入力=128、size=512のとき、SR=1, 2, 4のいずれの場合も出力=512で、内容はほとんど同じ（同一？）である。さらに、Keep=OFF、入力=128、Image size=256、SR=4の場合も出力=512で、内容は先ほどの3枚とほとんど同じ（同一？）である。

▼ 起動を簡単にする

　UIの起動が若干の手間なので、省力化することにします。

起動用バッチファイルの作成

　新規ファイル「C:\aiwork\diffb-g.bat」を下記の内容で作成してください。ファイル名は異なっていても構いません。

@echo off
cd \aiwork\DiffBIR
call venv\Scripts\activate.bat
python gradio_diffbir.py --ckpt ./general_full_v1.ckpt --config configs/model/cldm.yaml --reload_swinir --swinir_ckpt ./general_swinir_v1.ckpt
cd ..
deactivate

　これ以降にDiffBIRを起動したい場合は、コマンドプロンプトを開いて、下記のコマンドを順に実行してください。

cd \aiwork
diffb-g

　終了したい場合は、コマンドプロンプトのウインドウで「Ctrl + C」を押して、「バッチジョブを終了しますか (Y/N)」で「N」と答えるか、おもむろにウインドウを閉じてください。

　なお、このバッチファイルは「deactivate」が最後に書いてあり、これが実行されるとバッチファイルの処理が終了します。続けて処理を行いたい方は、この行を「call venv\Scripts\deactivate.bat」に変更してください。

▼ 顔の復元を行う

　先ほどは「general_*.ckpt」のモデルを使用しましたが、今度は「face_*.ckpt」のモデルを使用して顔の復元を行ってみます。

モデルのダウンロード

　https://huggingface.co/lxq007/DiffBIR/tree/main から「face_full_v1.ckpt」と「face_swinir_v1.ckpt」をダウンロードして、「C:\aiwork\DiffBIR」に移動してください。

起動用バッチファイルの作成

　新規ファイル「C:\aiwork\diffb-f.bat」を下記の内容で作成してください。ファイル名は異なっていても構いません。

@echo off
cd \aiwork\DiffBIR
call venv\Scripts\activate.bat
python gradio_diffbir.py --ckpt ./face_full_v1.ckpt --config configs/model/cldm.yaml --reload_swinir --swinir_ckpt ./face_swinir_v1.ckpt
cd ..
deactivate

起動

　コマンドプロンプトを開いて、下記のコマンドを順に実行してください。

cd \aiwork
diffb-f

　起動が完了すると、下記の表示で停止しますので、Webブラウザで http://127.0.0.1:7860/ へアクセスしてください。

Running on local URL:  http://0.0.0.0:7860

To create a public link, set `share=True` in `launch()`.

復元の実行

※後で実写ライクな画像に差し替えるかもしれません。

　512x512の画像（230908b_512.jpg）をダウンロードしてください。劣化させるため、低解像度で圧縮率を上げて保存し、それを読み込んで解像度を上げています。実写でも構いませんが、同じような手順で劣化したイメージを用意してください。解像度は256x256または512x512が良いでしょう。

　入力と出力の解像度等については、先に説明した「解像度等に関する説明（重要）」を参照してください。

▼ オプション

　復元に関するオプションは、下記の画像のとおりです。「SR Scale」、「Image size」、「Keep Original Size」については既に説明しました。

「Images」は一度に生成する画像の枚数です。2枚目以降は「Seed」の異なる画像と思われますが、詳細は不明です。
「Positive/Negative Prompt」「Prompt Guidance Scale」「Control Strength」「Steps」「Seed」は、Stable Diffusionの設定です。実は復元のキーとなる重要な設定なのですが、一言で説明するのが難しいので割愛します。
「Disable Preprocess Model」にチェックを入れると、Stage 1の処理（ノイズの除去）を行いません。Stage 1の効果、重要性を確認するためのものと思われます。
「Use Color Correction」は何らかの色補正と見られますが、効果は未確認です。

▼ メモ

pipのキャッシュを削除する

　pipコマンドでインストールしたパッケージは、ディスクにキャッシュされます。再利用時は便利ですが、例えばPyTorchは容量が2GB以上と大きいため、気になる方もいらっしゃるかと思います。
　パッケージのキャッシュを削除するには「pip cache purge」のコマンドを実行するか、「%LOCALAPPDATA%\pip\Cache」の中身を手動で削除してください。

アンインストールの手順

　もし、必要なファイルがあれば移動してください。その後、実行を終了（当該コマンドプロンプトのウインドウを閉じる）した状態で「DiffBIR」のディレクトリを削除してください。上記のキャッシュも削除すると良いでしょう。

　その他、「C:\Users\（ユーザー名）\.cache\torch\hub\checkpoints\」に「alexnet-owt-7be5be79.pth」がある場合は、DiffBIRがダウンロードしたものと思われます。こちらは削除が可能で、必要が生じた際は再度ダウンロードされます。

▼ その他

　私が書いた他の記事は、メニューよりたどってください。

　noteのアカウントはメインの@Mayu_Hiraizumiに紐付けていますが、記事に関することはサブアカウントの@riddi0908までお願いします。

この記事が気に入ったらサポートをしてみませんか？