nixified-ai/flake

trouble running out of GPU memory & setting env var

crabdancing opened this issue · 0 comments

When I try a basic prompt for invokeai-nvidia, this is the result:

Traceback (most recent call last):
  File "/nix/store/h4035a8gz1kqyvza3lpqw77qmrny9hzs-python3.10-InvokeAI-2.3.1.post2/lib/python3.10/site-packages/ldm/generate.py", line 557, in prompt2image
    results = generator.generate(
  File "/nix/store/h4035a8gz1kqyvza3lpqw77qmrny9hzs-python3.10-InvokeAI-2.3.1.post2/lib/python3.10/site-packages/ldm/invoke/generator/base.py", line 115, in generate
    image = make_image(x_T)
  File "/nix/store/h4035a8gz1kqyvza3lpqw77qmrny9hzs-python3.10-InvokeAI-2.3.1.post2/lib/python3.10/site-packages/ldm/invoke/generator/txt2img.py", line 45, in make_image
    pipeline_output = pipeline.image_from_embeddings(
  File "/nix/store/h4035a8gz1kqyvza3lpqw77qmrny9hzs-python3.10-InvokeAI-2.3.1.post2/lib/python3.10/site-packages/ldm/invoke/generator/diffusers_pipeline.py", line 429, in image_from_embeddings
    image = self.decode_latents(result_latents)
  File "/nix/store/h4035a8gz1kqyvza3lpqw77qmrny9hzs-python3.10-InvokeAI-2.3.1.post2/lib/python3.10/site-packages/ldm/invoke/generator/diffusers_pipeline.py", line 758, in decode_latents
    return super().decode_latents(latents)
  File "/nix/store/znwn5m13gh87v74wbvl2gfmyr2hykwqw-python3.10-diffusers-0.14.0/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py", line 426, in decode_latents
    image = self.vae.decode(latents).sample
  File "/nix/store/znwn5m13gh87v74wbvl2gfmyr2hykwqw-python3.10-diffusers-0.14.0/lib/python3.10/site-packages/diffusers/models/autoencoder_kl.py", line 185, in decode
    decoded = self._decode(z).sample
  File "/nix/store/znwn5m13gh87v74wbvl2gfmyr2hykwqw-python3.10-diffusers-0.14.0/lib/python3.10/site-packages/diffusers/models/autoencoder_kl.py", line 172, in _decode
    dec = self.decoder(z)
  File "/nix/store/k7f999ns4h0v0zb3yjnpka3935pydw2w-python3.10-torch-1.13.1/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
    return forward_call(*input, **kwargs)
  File "/nix/store/znwn5m13gh87v74wbvl2gfmyr2hykwqw-python3.10-diffusers-0.14.0/lib/python3.10/site-packages/diffusers/models/vae.py", line 188, in forward
    sample = up_block(sample)
  File "/nix/store/k7f999ns4h0v0zb3yjnpka3935pydw2w-python3.10-torch-1.13.1/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
    return forward_call(*input, **kwargs)
  File "/nix/store/znwn5m13gh87v74wbvl2gfmyr2hykwqw-python3.10-diffusers-0.14.0/lib/python3.10/site-packages/diffusers/models/unet_2d_blocks.py", line 1949, in forward
    hidden_states = upsampler(hidden_states)
  File "/nix/store/k7f999ns4h0v0zb3yjnpka3935pydw2w-python3.10-torch-1.13.1/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
    return forward_call(*input, **kwargs)
  File "/nix/store/znwn5m13gh87v74wbvl2gfmyr2hykwqw-python3.10-diffusers-0.14.0/lib/python3.10/site-packages/diffusers/models/resnet.py", line 131, in forward
    hidden_states = F.interpolate(hidden_states, scale_factor=2.0, mode="nearest")
  File "/nix/store/k7f999ns4h0v0zb3yjnpka3935pydw2w-python3.10-torch-1.13.1/lib/python3.10/site-packages/torch/nn/functional.py", line 3922, in interpolate
    return torch._C._nn.upsample_nearest2d(input, output_size, scale_factors)
torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 256.00 MiB (GPU 0; 5.80 GiB total capacity; 4.41 GiB already allocated; 251.38 MiB free; 4.58 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

What is the recommended way of setting environment variables? I tried setting them outside of the flake, i.e. PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:512 nix run .#invokeai-nvidia , but predictably, that violates reproduciblity guarantees and so doesn't work. I'm open to setting this any other way, if there's a working option I can toggle e.g. directly in the GUI or an imperative config file.