Stable Cascadeを触ってみつつ、サンプルコードを改良してVRAM消費を抑える

　StabilityAIから新しいText to ImageモデルStable Cascadeが発表されました。XLの正式公開が2023年7月26日なので半年しか経ってません。そして全く互換性がありません。XLのブーム始まったばかりなのに早いって。

　というわけで生成してみます。一応生成までの道のりを記しておきます。

1．ローカルへのインストール
コマンドプロンプトで適当なフォルダに赴き、git cloneを実行。gitが入ってない場合はgithubからzipをダウンロードして解凍でも大丈夫でしょう。

git clone https://github.com/Stability-AI/StableCascade

2．必須モジュールのインストール
CD で新しくできたフォルダに入り

pip install -r requirements.txt

しばし待つ。

3．モデルのダウンロード
以下は必須

effnet_encoder.safetensors
previewer.safetensors
stage_a.safetensors

Stage BとStage Cはお好みで。どの組み合わせでも使えるようだけれど･･････。とりあえず私はLiteをダウンロード。

stage_b_lite_bf16.safetensors
stage_c_lite_bf16.safetensors

4.設定ファイルの用意
　StableCascade\configs\inferenceにある、stage_b_3b.yamlを以下のように書き換えてstage_b_lite_3b.yamlとして保存する。

# GLOBAL STUFF
model_version: 700M
dtype: bfloat16

# For demonstration purposes in reconstruct_images.ipynb
webdataset_path: file:inference/imagenet_1024.tar
batch_size: 1
image_size: 1024
grad_accum_steps: 1

effnet_checkpoint_path: models/effnet_encoder.safetensors
stage_a_checkpoint_path: models/stage_a.safetensors
generator_checkpoint_path: models/stage_b_lite_bf16.safetensors

　stage_c_3b.yamlを以下のように書き換えてstage_c_lite_3b.yamlとして保存する。

# GLOBAL STUFF
model_version: 1B
dtype: bfloat16

effnet_checkpoint_path: models/effnet_encoder.safetensors
previewer_checkpoint_path: models/previewer.safetensors
generator_checkpoint_path: models/stage_c_lite_bf16.safetensors

5.推論用コードの準備
　サンプルコードを参考に画像を保存できるようにする。これをStableCascadeのルートフォルダに置いてGo!
なんか巨大なファイルをダウンロードした後推論が始まる。

import os
import yaml
import torch
import time
from tqdm import tqdm

#os.chdir('..')
from inference.utils import *
from core.utils import load_or_fail
from train import WurstCoreC, WurstCoreB

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(device)

# SETUP STAGE C
config_file = 'configs/inference/stage_c_lite_3b.yaml'
with open(config_file, "r", encoding="utf-8") as file:
    loaded_config = yaml.safe_load(file)

core = WurstCoreC(config_dict=loaded_config, device=device, training=False)

# SETUP STAGE B
config_file_b = 'configs/inference/stage_b_lite_3b.yaml'
with open(config_file_b, "r", encoding="utf-8") as file:
    config_file_b = yaml.safe_load(file)
    
core_b = WurstCoreB(config_dict=config_file_b, device=device, training=False)


# SETUP MODELS & DATA
extras = core.setup_extras_pre()
models = core.setup_models(extras)
models.generator.eval().requires_grad_(False)
print("STAGE C READY")

extras_b = core_b.setup_extras_pre()
models_b = core_b.setup_models(extras_b, skip_clip=True)
models_b = WurstCoreB.Models(
   **{**models_b.to_dict(), 'tokenizer': models.tokenizer, 'text_model': models.text_model}
)
models_b.generator.bfloat16().eval().requires_grad_(False)
print("STAGE B READY")

batch_size = 1
# caption = "Cinematic photo of an anthropomorphic nerdy rodent sitting in a cafe reading a book"
caption = "Cinematic photo of a girl in cafe"
height, width = 832,1216
stage_c_latent_shape, stage_b_latent_shape = calculate_latent_sizes(height, width, batch_size=batch_size)

# Stage C Parameters
extras.sampling_configs['cfg'] = 4
extras.sampling_configs['shift'] = 2
extras.sampling_configs['timesteps'] = 20
extras.sampling_configs['t_start'] = 1.0

# Stage B Parameters
extras_b.sampling_configs['cfg'] = 1.1
extras_b.sampling_configs['shift'] = 1
extras_b.sampling_configs['timesteps'] = 10
extras_b.sampling_configs['t_start'] = 1.0

# PREPARE CONDITIONS
batch = {'captions': [caption] * batch_size}
conditions = core.get_conditions(batch, models, extras, is_eval=True, is_unconditional=False, eval_image_embeds=False)
unconditions = core.get_conditions(batch, models, extras, is_eval=True, is_unconditional=True, eval_image_embeds=False)    
conditions_b = core_b.get_conditions(batch, models_b, extras_b, is_eval=True, is_unconditional=False)
unconditions_b = core_b.get_conditions(batch, models_b, extras_b, is_eval=True, is_unconditional=True)

with torch.no_grad(), torch.cuda.amp.autocast(dtype=torch.bfloat16):
    # torch.manual_seed(2244096)

    sampling_c = extras.gdf.sample(
        models.generator, conditions, stage_c_latent_shape,
        unconditions, device=device, **extras.sampling_configs,
    )
    for (sampled_c, _, _) in tqdm(sampling_c, total=extras.sampling_configs['timesteps']):
        sampled_c = sampled_c
        
    # preview_c = models.previewer(sampled_c).float()
    # show_images(preview_c)

    conditions_b['effnet'] = sampled_c
    unconditions_b['effnet'] = torch.zeros_like(sampled_c)

    sampling_b = extras_b.gdf.sample(
        models_b.generator, conditions_b, stage_b_latent_shape,
        unconditions_b, device=device, **extras_b.sampling_configs
    )
    for (sampled_b, _, _) in tqdm(sampling_b, total=extras_b.sampling_configs['timesteps']):
        sampled_b = sampled_b
    sampled = models_b.stage_a.decode(sampled_b).float()

time = time.strftime(r"%Y%m%d%H%M%S")
for i, img in enumerate(sampled):
    img = torchvision.transforms.functional.to_pil_image(img.clamp(0, 1))
    fileName = f"img-{time}-{i+1}.png"
    img.save(fileName)

　完了。StableCascadeフォルダに生成されました。VRAM消費は6G、生成時間は30秒程度。

Cinematic photo of a girl in cafe

　どうにも品質がよろしくないようなので、Liteじゃなくふつうのモデルを使ってみます。stage_c_3b.yamlとstage_b_3b.yamlを使います。

　あれ。VRAMが足りません。VRAM12Gでは共有領域にはみ出してしまいました。うむむ。きびしい。ここで泣いてふて寝するわけにはいきません。工夫してみましょう。原理的にはStageCとStageBの間にはLatentの受け渡ししか発生しないはずですから、StageCの推論が終わった後、StageCをunloadしてStageBをロードすればVRAMを削減できるはずです。というかサンプルコードはいきなり両方のモデルをロードするところから始まるとか石油王仕様なんですかね･･････。

import os
import yaml
import torch
import time
from tqdm import tqdm
from transformers import AutoTokenizer, CLIPTextModelWithProjection
#os.chdir('..')
from inference.utils import *
from core.utils import load_or_fail
from train import WurstCoreC, WurstCoreB

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(device)

# SETUP STAGE C
config_file = 'configs/inference/stage_c_3b.yaml'
with open(config_file, "r", encoding="utf-8") as file:
    loaded_config = yaml.safe_load(file)

core = WurstCoreC(config_dict=loaded_config, device=device, training=False)

# SETUP STAGE B
config_file_b = 'configs/inference/stage_b_3b.yaml'
with open(config_file_b, "r", encoding="utf-8") as file:
    config_file_b = yaml.safe_load(file)
    
core_b = WurstCoreB(config_dict=config_file_b, device=device, training=False)

# SETUP MODELS & DATA
extras = core.setup_extras_pre()
models = core.setup_models(extras)
models.generator.eval().requires_grad_(False)
print("STAGE C READY")

batch_size = 1
# caption = "Cinematic photo of an anthropomorphic nerdy rodent sitting in a cafe reading a book"
caption = "Cinematic photo of a girl in cafe"
height, width = 832,1216
stage_c_latent_shape, stage_b_latent_shape = calculate_latent_sizes(height, width, batch_size=batch_size)

# Stage C Parameters
extras.sampling_configs['cfg'] = 4
extras.sampling_configs['shift'] = 2
extras.sampling_configs['timesteps'] = 20
extras.sampling_configs['t_start'] = 1.0

# PREPARE CONDITIONS
batch = {'captions': [caption] * batch_size}
conditions = core.get_conditions(batch, models, extras, is_eval=True, is_unconditional=False, eval_image_embeds=False)
unconditions = core.get_conditions(batch, models, extras, is_eval=True, is_unconditional=True, eval_image_embeds=False)    

with torch.no_grad(), torch.cuda.amp.autocast(dtype=torch.bfloat16):
    # torch.manual_seed(2244096)

    sampling_c = extras.gdf.sample(
        models.generator, conditions, stage_c_latent_shape,
        unconditions, device=device, **extras.sampling_configs,
    )
    for (sampled_c, _, _) in tqdm(sampling_c, total=extras.sampling_configs['timesteps']):
        sampled_c = sampled_c
        
    # preview_c = models.previewer(sampled_c).float()
    # show_images(preview_c)
del models
import gc
gc.collect()
dtype = getattr(torch, core.config.dtype) if core.config.dtype else torch.float32
tokenizer = AutoTokenizer.from_pretrained(core.config.clip_text_model_name)
text_model = CLIPTextModelWithProjection.from_pretrained(core.config.clip_text_model_name).requires_grad_(False).to(dtype).to(core.device)

extras_b = core_b.setup_extras_pre()
models_b = core_b.setup_models(extras_b, skip_clip=True)
models_b = WurstCoreB.Models(
   **{**models_b.to_dict(), 'tokenizer': tokenizer, 'text_model': text_model}
)
models_b.generator.bfloat16().eval().requires_grad_(False)
print("STAGE B READY")

# Stage B Parameters
extras_b.sampling_configs['cfg'] = 1.1
extras_b.sampling_configs['shift'] = 1
extras_b.sampling_configs['timesteps'] = 10
extras_b.sampling_configs['t_start'] = 1.0
conditions_b = core_b.get_conditions(batch, models_b, extras_b, is_eval=True, is_unconditional=False)
unconditions_b = core_b.get_conditions(batch, models_b, extras_b, is_eval=True, is_unconditional=True)

with torch.no_grad(), torch.cuda.amp.autocast(dtype=torch.bfloat16):
    conditions_b['effnet'] = sampled_c
    unconditions_b['effnet'] = torch.zeros_like(sampled_c)

    sampling_b = extras_b.gdf.sample(
        models_b.generator, conditions_b, stage_b_latent_shape,
        unconditions_b, device=device, **extras_b.sampling_configs
    )
    for (sampled_b, _, _) in tqdm(sampling_b, total=extras_b.sampling_configs['timesteps']):
        sampled_b = sampled_b
    sampled = models_b.stage_a.decode(sampled_b).float()

time = time.strftime(r"%Y%m%d%H%M%S")
for i, img in enumerate(sampled):
    img = torchvision.transforms.functional.to_pil_image(img.clamp(0, 1))
    fileName = f"img-{time}-{i+1}.png"
    img.save(fileName)

　単純にdelするだけではVRAM消費が消えませんでした。どうやらStage_cの要素をStageBに割り当てているため元のコードだとStageCをunloadできないみたい。そんなことする必要があるとは思えないのですが、わざわざunloadを難しくしているように思えます。というわけで、Stage_bでも使うtokenizerとtext_modelだけ別途ロードすることにする。

dtype = getattr(torch, core.config.dtype) if core.config.dtype else torch.float32
tokenizer = AutoTokenizer.from_pretrained(core.config.clip_text_model_name)
text_model = CLIPTextModelWithProjection.from_pretrained(core.config.clip_text_model_name).requires_grad_(False).to(dtype).to(core.device)

　これで動かすとStageCでは9G程度、その後unlaodが行われ、StageBでは5G程度のVRAM消費で動きました。めでたしめでたし。推論時間は2分弱と大幅に伸びました。

Cinematic photo of a girl in cafe

　こちらの方がきれいに出てますね。3.6Bぐらいないときれいなのは出ないようです。適当なプロンプト、そのままの条件で出したにしてはとってもきれいですね。

　そろそろ眠いですが最後にイラスト系も出してみますか。

ほう。これはなかなか

　なかなかいい感じですね。danbooruは学習していないようなので、学習次第ではかなりの性能を秘めているように感じました。XLより重くて時間が掛かるのがネックですが、そこら辺の高速化技術もすぐに出てくるでしょう。そろそろ眠いのでねることにします。おやすみなさい。