support Waifu Diffusion on Google Colab Free

README.md

Google Colab 無料枠で Waifu Diffusion を動かす

概要

• Waifu Diffusion を Colab 無料枠で動かそうとするとメモリが溢れてしまう
  • wd-v1-2-full-ema.cpkt には画像生成するだけなら不要な情報が乗っているため
• 本来ならfull-emaでないモデルを用意すべき
• 暫定案としてVRAMに直接ロードするようパッチを当てると、メモリが節約できてギリギリ動作する

対応

scripts/txt2img.py, scripts/img2img.py に上記 diff を当てて cuda利用可能環境ではcpuを介さずgpuに直接ロードする

img.py.diff

 def load_model_from_config(config, ckpt, verbose=False):
     print(f"Loading model from {ckpt}")
-    pl_sd = torch.load(ckpt, map_location="cpu")
+    device = torch.device(
+        "cuda") if torch.cuda.is_available() else torch.device("cpu")
+    pl_sd = torch.load(ckpt, map_location=device)