Pythonライブラリ (超解像|画質向上)：DiffBIR
１．緒言
低画質の画像を高画質に変える技術である”超解像”として「DiffBIR」を紹介します。結論として、GPUでの実装まではできなかったため、CPUで時間かけても良い人向けとなります。
ご参考までに同様の技術としてReal-ESRGANも紹介していますのでご確認ください。
１－１．DiffBIRの概要
下図の通り、ぼやけた画像を鮮明にしたりできます。他の技術との優位性比較をしており、見たところ高性能のように見えます。
１－２．DiffBIRのアーキテクチャ
DiffBIRのアーキテクチャは下記の通りです。詳細は理解できませんが、2段のパイプラインで処理しており、Stable DiffusionやVAEエンコーダを使用しており、GANとは異なっていそうです。
２．環境構築：Windows編
Windowsで環境構築する場合を記載します。
２－１．Windowsとの相性確認
大前提として公式はLinux OS以外の使用を推奨していません。
Windowsの場合は「tritonのインストール時に躓くため、GPUでなくCPUで動かせる。CUDAを使うならissue#24で確認して」とあります。
よって環境構築の方針は下記２点で対応しました。
CPUのみで使用できるようにする
頑張ってGPUも使えるようにする（tritonをインストールする
２－２．仮想環境の作成
仮想環境をCondaで作成します。ここでのポイントは「公式はPython3.9を推奨しているが、GPUを使いたいならPython3.10を使用」です。
環境が出来たら仮想環境を有効化します。
[Terminal:CPUだけでよい人]
conda create -n diffbir python=3.9
conda activate diffbir[Terminal：GPU使いたい人]
conda create -n diffbir python=3.10
conda activate diffbir２－３．ライブラリインストール
ライブラリをインストールします。前述の通りWindowsではそのままだとtritonを入れられないためエラーが発生しますので少し対応します。
２－３－１．CPUのみ使用用
GPUを使わない人は”requirements.txt”を開いてtritonをコメントアウトしてから下記コマンド実行します。
[Terminal]
pip install -r requirements.txt
【参考：コメントアウトしない場合】
tritonのコメントアウトをしない場合は、公式が説明している通りエラーが発生しました。
２－３－２．GPUで実行用
下記issueにHuggingfaceが提供するカスタムビルドされたTritonのバイナリが紹介されています。このバイナリを使用してTritonをインストールします。
注意点としてPython3.9だと下記のようなエラーが発生します。必ず仮想環境構築時はPython3.10系を使用してください。
実装方法としてはtritonをコメントアウトした後にpip installを実行して、次に別途tritonをインストールしました。
[Terminal]
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
私の環境では問題なくインストールできました。
２－４．学習済みモデルの重みを取得
weightsフォルダを作成して、公式の下記からダウンロードしたファイルをweightsフォルダ内に保存します。
３．環境構築：Google Colab編
推奨はLinux OSのみのためGoogle Colabを使用していきます。
現状では環境構築が出来なかったため、うまくいったら記事を修正します。参考までにPythonのVersionは3.10系です。
【パターン１：公式通り】
Google Colab内で仮想環境を作るのは難しいため、仮想環境の作成を除いて公式通りに進めていきます。
まずはレポジトリをクローンしてライブラリをインストールします。
[Terminal]
!git clone https://github.com/XPixelGroup/DiffBIR.git
%cd DiffBIR
!pip install -r requirements.txt次にweightsフォルダを作成して、その中に必要な学習済みモデルをインストールします。
[Terminal]
!mkdir weights
!wget -O weights/general_full_v1.ckpt https://huggingface.co/lxq007/DiffBIR/resolve/main/general_full_v1.ckpt
!wget -O weights/general_swinir_v1.ckpt https://huggingface.co/lxq007/DiffBIR/resolve/main/general_swinir_v1.ckpt 最後にコードを実行します。するとPytorchに関するエラーが発生しました。torchdataと記載がありますがおそらくpytorch lightning関連の互換性のエラーだと予想されます。
修復がムリゲーだったためあきらめました。
[Terminal]
!python inference.py --input inputs/demo/general --output results/demo/general --config configs/model/cldm.yaml --ckpt weights/general_full_v1.ckpt --device cuda[OUT]
Traceback (most recent call last):
File "/content/DiffBIR/DiffBIR/DiffBIR/inference.py", line 9, in <module>
import pytorch_lightning as pl
File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/__init__.py", line 20, in <module>
from pytorch_lightning import metrics # noqa: E402
File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/metrics/__init__.py", line 15, in <module>
from pytorch_lightning.metrics.classification import ( # noqa: F401
File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/metrics/classification/__init__.py", line 14, in <module>
from pytorch_lightning.metrics.classification.accuracy import Accuracy # noqa: F401
File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/metrics/classification/accuracy.py", line 18, in <module>
from pytorch_lightning.metrics.utils import deprecated_metrics, void
File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/metrics/utils.py", line 29, in <module>
from pytorch_lightning.utilities import rank_zero_deprecation
File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/utilities/__init__.py", line 18, in <module>
from pytorch_lightning.utilities.apply_func import move_data_to_device # noqa: F401
File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/utilities/apply_func.py", line 34, in <module>
from torchtext.legacy.data import Batch
File "/usr/local/lib/python3.10/dist-packages/torchtext/__init__.py", line 12, in <module>
from . import data, datasets, prototype, functional, models, nn, transforms, utils, vocab, experimental
File "/usr/local/lib/python3.10/dist-packages/torchtext/datasets/__init__.py", line 3, in <module>
from .ag_news import AG_NEWS
File "/usr/local/lib/python3.10/dist-packages/torchtext/datasets/ag_news.py", line 12, in <module>
from torchdata.datapipes.iter import FileOpener, IterableWrapper
File "/usr/local/lib/python3.10/dist-packages/torchdata/__init__.py", line 9, in <module>
from . import datapipes
File "/usr/local/lib/python3.10/dist-packages/torchdata/datapipes/__init__.py", line 9, in <module>
from . import iter, map, utils
File "/usr/local/lib/python3.10/dist-packages/torchdata/datapipes/iter/__init__.py", line 124, in <module>
from torchdata.datapipes.iter.util.sharding import (
File "/usr/local/lib/python3.10/dist-packages/torchdata/datapipes/iter/util/sharding.py", line 9, in <module>
from torch.utils.data.datapipes.iter.sharding import SHARDING_PRIORITIES
ModuleNotFoundError: No module named 'torch.utils.data.datapipes.iter.sharding'【パターン２：公式のColab】
公式GitHubから”Open in Colab”を開いてコマンドを実行しました。
[DiffBIR_colab.ipynb]
%cd /content
!git clone -b dev https://github.com/camenduru/DiffBIR
%cd /content/DiffBIR
!pip install -q einops pytorch_lightning gradio omegaconf transformers lpips # opencv-python
!pip install -q https://download.pytorch.org/whl/cu118/xformers-0.0.22.post4%2Bcu118-cp310-cp310-manylinux2014_x86_64.whl
!pip install -q git+https://github.com/mlfoundations/open_clip@v2.20.0
!apt -y install -qq aria2
# !aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/lxq007/DiffBIR/resolve/main/face_full_v1.ckpt -d /content/models -o face_full_v1.ckpt
# !aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/lxq007/DiffBIR/resolve/main/face_swinir_v1.ckpt -d /content/models -o face_swinir_v1.ckpt
!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/lxq007/DiffBIR/resolve/main/general_full_v1.ckpt -d /content/models -o general_full_v1.ckpt
!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/lxq007/DiffBIR/resolve/main/general_swinir_v1.ckpt -d /content/models -o general_swinir_v1.ckpt
# !python gradio_diffbir.py --ckpt /content/models/face_full_v1.ckpt --config /content/DiffBIR/configs/model/cldm.yaml --reload_swinir --swinir_ckpt /content/models/face_swinir_v1.ckpt
!python gradio_diffbir.py --ckpt /content/models/general_full_v1.ckpt --config /content/DiffBIR/configs/model/cldm.yaml --reload_swinir --swinir_ckpt /content/models/general_swinir_v1.ckptinputsに画像を手動で入れて、モデルの場所を移動させても指定したフォルダ(--output results)に結果が出力されません。エラー文も出ないため原因不明です。
[Terminal]
import os
if not os.path.exists('inputs'):
os.mkdir('inputs')
!mkdir weights
!wget -O weights/general_full_v1.ckpt https://huggingface.co/lxq007/DiffBIR/resolve/main/general_full_v1.ckpt
!wget -O weights/general_swinir_v1.ckpt https://huggingface.co/lxq007/DiffBIR/resolve/main/general_swinir_v1.ckpt
!python inference.py --input inputs --config configs/model/cldm.yaml --ckpt weights/general_full_v1.ckpt --reload_swinir --swinir_ckpt weights/general_swinir_v1.ckpt --steps 50 --sr_scale 4 --output results 
４．実行編：Windows
それでは実際に実行して超解像を試してみます。
４－１．CPUのみ：--device cpu
コマンドの"--device"オプションをCPUに指定します。その他のオプションは下記の通りです（inference.pyのdef parse_args()参照）。
--input inputs/demo/general： 入力画像が保存ディレクトリを指定
この場合”inputs/demo/general”ディレクトリ内の画像を処理
--config configs/model/cldm.yaml：モデルの設定ファイルのパスを指定
--ckpt weights/general_full_v1.ckpt：モデルの重みのファイルパスを指定
general_full_v1.ckpt は公式GitHubからダウンロードしてweightsに保存
--reload_swinir： SwinIRモデルの重みをリロード
--swinir_ckpt weights/general_swinir_v1.ckpt： SwinIRモデル重みのパス
general_swinir_v1.ckpt は公式からダウンロードしてweightsに保存。
--steps 50： 処理の際に実行するステップ数を指定
--sr_scale 4：高解像度化のスケールを指定（画像のサイズを4倍に拡大）
--color_fix_type wavelet：色補正のタイプを指定
--output results/demo/general：出力結果を保存するディレクトリを指定
--device cuda：処理に使用するデバイスを指定
--tiled, --tile_size 512, --tile_stride 256：画像をタイルに分割して処理？
実施したけどエラーが発生（※原因不明
[Terminal]
python inference.py \
--input inputs/demo/general \
--config configs/model/cldm.yaml \
--ckpt weights/general_full_v1.ckpt \
--reload_swinir --swinir_ckpt weights/general_swinir_v1.ckpt \
--steps 50 \
--sr_scale 4 \
--color_fix_type wavelet \
--output results/demo/general \
--device cpu[Terminal※1列Ver.]
python inference.py --input inputs/demo/general --config configs/model/cldm.yaml --ckpt weights/general_full_v1.ckpt --reload_swinir --swinir_ckpt weights/general_swinir_v1.ckpt --steps 50 --sr_scale 4 --color_fix_type wavelet --output results/demo/general --device cpu
問題なく実行されると出力先が”save to <保存ディレクトリ>”で表示（--outputオプションで指定した場所と同じ）されており、画像名は接尾語として"_0"がついております。
CPU（AMD Ryzen 7 3700X 8-Core Processor 3.59 GHz）だけだと、1枚の画像で1hくらいかかりました。
【結果の確認】
結果は下記の通りです。
４－２．GPUで処理
GPUで処理したい場合は"--device cuda"とします。デフォルトはcudaのため、おそらく記載なしでもいけると思います。
[Terminal]
python inference.py --input inputs/demo/general --config configs/model/cldm.yaml --ckpt weights/general_full_v1.ckpt --reload_swinir --swinir_ckpt weights/general_swinir_v1.ckpt --steps 50 --sr_scale 4 --color_fix_type wavelet --output results/demo/general --device cuda[OUT]
torch.cuda.OutOfMemoryError: CUDA out of memory.
Tried to allocate 468.00 MiB (GPU 0; 8.00 GiB total capacity; 6.94 GiB already allocated;
0 bytes free; 7.01 GiB reserved in total by PyTorch)
If reserved memory is >> allocated memory try setting max_split_size_mb to avoid
fragmentation. See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF私のPCはNVIDIA GeForce RTX 2060(8GB)を積んでますが、メモリが全然足りず動きませんでした。
参考資料
参考資料１：実装
参考資料２：技術
参考資料３：その他
あとがき
そろそろちゃんと実装したい・・・・
あと、メモリの大きいGPUが無いとなんもできないけど金がねえ・・・
この記事が気に入ったらサポートをしてみませんか？