見出し画像

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

まゆひらa

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 


▼ インストール

　リポジトリの説明は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が良いでしょう。

230908a_512（こちらはダウンロードしないでください）

　次に、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が良いでしょう。

230908b_512（こちらはダウンロードしないでください）

　次に、UIの左側の欄へ画像をドラッグ＆ドロップしてから「Run」をクリックしてください。256x256の画像の場合は、「SR Scale」を2に設定すると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までお願いします。

　

　

いいなと思ったら応援しよう！