Skip to content

Instantly share code, notes, and snippets.

@td2sk
Created September 8, 2022 14:31
Show Gist options
  • Save td2sk/e32a39344537fb3cd756ef4abdd3d371 to your computer and use it in GitHub Desktop.
Save td2sk/e32a39344537fb3cd756ef4abdd3d371 to your computer and use it in GitHub Desktop.
support Waifu Diffusion on Google Colab Free

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に直接ロードする

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)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment