Ubuntu22.04DesktopでSeeSRをインストールした記録

画像の超解像ツールのSeeSRのインストールしたので手順を残します。

GitHub - cswry/SeeSR: [CVPR2024] SeeSR: Towards Semantics-Aware Real-World Image Super-Resolution

SeeSR: Towards Semantics-Aware Real-World Image Super-Resolution - cswry/SeeSR

実行した環境

Windows11のWSL2でUbuntu22.04

CUDAのバージョンは11.8

githubからリポジトリをクローン

cd ~/git
git clone https://github.com/cswry/SeeSR.git
cd SeeSR

venvによるpipのインストール環境の構築

mkdir venv
python3 -m venv venv
source venv/bin/activate

requirements.txtの内容をテキストエディタで編集

torch==2.0.1 --index-url https://download.pytorch.org/whl/cu118
huggingface_hub==0.25.2
diffusers==0.21.0
pytorch_lightning==2.1.3
accelerate==0.25.0
transformers==4.25.0
xformers
loralib==0.1.2
fairscale==0.4.13
pydantic==1.10.11
gradio==3.24.0
opencv-python==4.9.0.80
chardet==5.2.0
einops==0.7.0
scipy==1.10.1
timm==0.9.12
gdown

torchはcudaのインストールバージョンに合わせる。

編集後インストール

pip install -r requirements.txt

model用ディレクトリ作成とStableDiffusionのインストールスクリプトを作成

import os
from diffusers import StableDiffusionPipeline, EulerDiscreteScheduler
import torch

os.makedirs(os.path.join("preset", "models"), exist_ok=True)
model_id = "stabilityai/stable-diffusion-2-base"
scheduler = EulerDiscreteScheduler.from_pretrained(model_id, subfolder="scheduler")
pipe = StableDiffusionPipeline.from_pretrained(model_id, scheduler=scheduler, torch_dtype=torch.float16)
save_directory = "preset/models/stable-diffusion-2-base"
pipe.save_pretrained(save_directory)
scheduler.save_pretrained(save_directory + "/scheduler")

適当なスクリプト名で作成し実行します。

preset/modelsというディレクトリが出来上がりそちらにmodelファイルを用意していきます。

Googleドライブで公開されているmodelファイルをダウンロードするスクリプトを作成

import os
import gdown

os.makedirs('preset/models/seesr', exist_ok=True)
os.makedirs('preset/models/seesr/unet', exist_ok=True)
os.makedirs('preset/models/seesr/controlnet', exist_ok=True)


# DAPE
gdown.download(
    'https://drive.google.com/uc?id=1KIV6VewwO2eDC9g4Gcvgm-a0LDI7Lmwm',
    'preset/models/DAPE.pth',
    quiet=False,
)

# SeeSR
gdown.download(
    'https://drive.google.com/uc?id=1sUU-noBhI9tGF9j6S8pIjO_AnmT99NDh',
    'preset/models/seesr/scaler.pt',
    quiet=False,
)
gdown.download(
    'https://drive.google.com/uc?id=1JxX_-MHZZ77Z0ElnEnoLhXi6MlTl7SqO',
    'preset/models/seesr/unet/config.json',
    quiet=False,
)
gdown.download(
    'https://drive.google.com/uc?id=10KEt-EYj565VhA8pJOF9Bxp_2Wdxo9lV',
    'preset/models/seesr/unet/diffusion_pytorch_model.safetensors',
    quiet=False,
)
gdown.download(
    'https://drive.google.com/uc?id=1hKNsEB9k2GC3lsBKJf1zA6Mi4b-tTY5Q',
    'preset/models/seesr/controlnet/config.json',
    quiet=False,
)
gdown.download(
    'https://drive.google.com/uc?id=1R_rWUdJ6oJXR6PhPlbGAvfEltajr1KbN',
    'preset/models/seesr/controlnet/diffusion_pytorch_model.safetensors',
    quiet=False,
)

こちらもスクリプトファイルを作成し実行します。

次にhuggingfaceで公開されているmodelファイルをwgetでダウンロードします。

wget https://huggingface.co/spaces/xinyu1205/recognize-anything/resolve/main/ram_swin_large_14m.pth -P preset/models

プリセット画像を拡大しています。

python3 test_seesr.py \
--pretrained_model_path preset/models/stable-diffusion-2-base \
--prompt '' \
--seesr_model_path preset/models/seesr \
--ram_ft_path preset/models/DAPE.pth \
--image_path preset/datasets/test_datasets \
--output_dir preset/datasets/output \
--start_point lr \
--num_inference_steps 50 \
--guidance_scale 5.5 \
--process_size 512

成功するとpreset/datasets/output/sample00/179.pngに画像ファイルが出来上がります。

メキシコサラマンダーの画像が128×128が512×512へ4倍に拡大されました。

modelsディレクトリの構造

preset/models
├── DAPE.pth
├── ram_swin_large_14m.pth
├── seesr
│   ├── controlnet
│   │   ├── config.json
│   │   └── diffusion_pytorch_model.safetensors
│   ├── scaler.pt
│   └── unet
│       ├── config.json
│       └── diffusion_pytorch_model.safetensors
└── stable-diffusion-2-base
    ├── feature_extractor
    │   └── preprocessor_config.json
    ├── model_index.json
    ├── scheduler
    │   └── scheduler_config.json
    ├── text_encoder
    │   ├── config.json
    │   └── model.safetensors
    ├── tokenizer
    │   ├── merges.txt
    │   ├── special_tokens_map.json
    │   ├── tokenizer_config.json
    │   └── vocab.json
    ├── unet
    │   ├── config.json
    │   └── diffusion_pytorch_model.safetensors
    └── vae
        ├── config.json
        └── diffusion_pytorch_model.safetensors

10 directories, 20 files

入力・出力ディレクトリを固定してバッチ処理を行うスクリプト

#!/bin/bash
cd ~
venvFile=`pwd`'/git/SeeSR/venv/bin/activate'
dataRoot='/mnt/f' # 環境に合わせて修正(Windowsのfドライブ)
inputDir=$dataRoot'/input' # 環境に合わせて修正（入力画像)
outputDir=$dataRoot'/output' # 環境に合わせて修正（出力画像)
logfile=$dataRoot'/logs/SeeSR.log' # 環境に合わせて修正（ログのパス)

cd $outputDir
if [ `ls -U1 | wc -l` -ne 0 ]; then
    rm *
fi;
cd $inputDir
if [ `ls -U1 | wc -l` -eq 0 ]; then
    echo 'Not File';
    exit 1
fi;
cd ~/git/SeeSR
source $venvFile
echo '*** SeeSR Start ***' >> $logfile
date >> $logfile
python3 test_seesr.py \
    --pretrained_model_path preset/models/stable-diffusion-2-base \
    --prompt '' \
    --seesr_model_path preset/models/seesr \
    --ram_ft_path preset/models/DAPE.pth \
    --image_path $inputDir \
    --output_dir $outputDir \
    --start_point lr \
    --num_inference_steps 50 \
    --guidance_scale 5.5 \
    --upscale 4 \
    --process_size 512
date >> $logfile
echo '*** SeeSR End ***' >> $logfile
deactivate
cd $inputDir
if [ `ls -U1 | wc -l` -ne 0 ]; then
    rm *
fi;

プリセット画像が128×128を4倍にする程度であれば、それほど時間がかかるわけではないですが、1100×1600の画像を拡大しようとすると数十分ほど時間がかかります。

PCはRyzen7-5700X MEM:64GB RTX4070といった構成です。

また、完走せずにkillされることが結構の頻度で発生して安定しない感じです。

実行中はVRAMで使用率が高まり、共有メモリとしてメインメモリーも使われて、さらにCドライブへのアクセスが増大しますので、かなりメモリを消費しています。

内部でStableDiffusion使われていることから、画像の内容を認識し、文字としてプロンプトにセットして画像を生成しているようです。そちらが超解像とどのような使われ方をしているのでしょうか？詳細のクオリティが高い部分があり、オリジナルには存在していないので、生成しているような感じがします。

SeeSRに限らずですが、実行環境の構築が難しく、そのため、このような記事で手順を残していますが、パッケージ、ドライバ、ツールキットなどがバージョンアップする関係上、再現性についてはよろしくない感じです。一度作成した実行環境を維持する仮想環境が必須だと思います。

pipパッケージレベル…venv

python(バージョン)…Anacondaなど

コンテナ(アプリケーション+実行環境)…docker

OSレベル…KVMなど

コメント