메모리 사용량 줄이기

Flux와 Wan과 같은 최신 diffusion 모델은 수십억 개의 파라미터를 가지고 있어 추론 시 하드웨어에서 상당한 메모리를 차지합니다. 이는 일반적인 GPU가 충분한 메모리를 갖추지 못하는 경우가 많아 도전적인 과제입니다. 이러한 메모리 제한을 극복하기 위해, 사용 가능한 경우 여러 GPU를 사용하거나 파이프라인의 일부 구성 요소를 CPU로 offload하는 등의 방법을 활용할 수 있습니다.

이 가이드에서는 메모리 사용량을 줄이는 방법을 보여 드립니다.

모델에 따라 이러한 기법을 조정해야 할 수 있음을 유의하세요. 예를 들어, Transformer 기반의 diffusion 모델은 UNet 기반 모델과 달리 이러한 메모리 최적화에서 동일한 이점을 얻지 못할 수 있습니다.

여러 GPU

여러 개의 GPU에 접근할 수 있다면, 대규모 모델을 하드웨어에 효율적으로 로드하고 분산하기 위한 몇 가지 옵션이 있습니다. 이러한 기능은 Accelerate 라이브러리가 지원하므로, 먼저 설치되어 있는지 확인하세요.

pip install -U accelerate

Sharded checkpoints

여러 개의 샤드로 큰 체크포인트를 로드하는 것은 유용합니다. 샤드가 하나씩 로드되기 때문에 메모리 사용량이 낮게 유지되며, 모델 크기와 가장 큰 샤드 크기만큼의 메모리만 필요하기 때문입니다. fp32 체크포인트가 5GB를 초과하는 경우 샤딩을 권장합니다. 기본 샤드 크기는 5GB입니다.

max_shard_size 매개변수를 사용하여 save_pretrained()에서 체크포인트를 샤드합니다.

from diffusers import AutoModel

unet = AutoModel.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", subfolder="unet"
)
unet.save_pretrained("sdxl-unet-sharded", max_shard_size="5GB")

sharded checkpoint를 사용하여 regular checkpoint 대신 메모리를 절약할 수 있습니다.

import torch
from diffusers import AutoModel, StableDiffusionXLPipeline

unet = AutoModel.from_pretrained(
    "username/sdxl-unet-sharded", dtype=torch.float16
)
pipeline = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    unet=unet,
    dtype=torch.float16
).to("cuda")  # or "mps", "xpu", "cpu"

장치 배치

Device placement는 실험적 기능이며 API가 변경될 수 있습니다. 현재는 balanced 전략만 지원됩니다. 향후 추가 매핑 전략을 지원할 계획입니다.

device_map 매개변수는 파이프라인의 모델 구성 요소 또는 개별 모델의 레이어가 여러 디바이스에 어떻게 분배되는지를 제어합니다.

파이프라인 수준

모델 수준

device_map을 직접 설계할 때는, 모델의 특정 모듈 이름 또는 레이어와 디바이스 식별자(GPU의 경우 정수, CPU의 경우 cpu, 디스크의 경우 disk)를 키-값 쌍으로 갖는 사전(dictionary) 형태여야 합니다.

모델에 hf_device_map을 호출하여 모델 레이어가 어떻게 분배되는지 확인한 후, 직접 설계해 보세요.

print(transformer.hf_device_map)
{'pos_embed': 0, 'time_text_embed': 0, 'context_embedder': 0, 'x_embedder': 0, 'transformer_blocks': 0, 'single_transformer_blocks.0': 0, 'single_transformer_blocks.1': 0, 'single_transformer_blocks.2': 0, 'single_transformer_blocks.3': 0, 'single_transformer_blocks.4': 0, 'single_transformer_blocks.5': 0, 'single_transformer_blocks.6': 0, 'single_transformer_blocks.7': 0, 'single_transformer_blocks.8': 0, 'single_transformer_blocks.9': 0, 'single_transformer_blocks.10': 'cpu', 'single_transformer_blocks.11': 'cpu', 'single_transformer_blocks.12': 'cpu', 'single_transformer_blocks.13': 'cpu', 'single_transformer_blocks.14': 'cpu', 'single_transformer_blocks.15': 'cpu', 'single_transformer_blocks.16': 'cpu', 'single_transformer_blocks.17': 'cpu', 'single_transformer_blocks.18': 'cpu', 'single_transformer_blocks.19': 'cpu', 'single_transformer_blocks.20': 'cpu', 'single_transformer_blocks.21': 'cpu', 'single_transformer_blocks.22': 'cpu', 'single_transformer_blocks.23': 'cpu', 'single_transformer_blocks.24': 'cpu', 'single_transformer_blocks.25': 'cpu', 'single_transformer_blocks.26': 'cpu', 'single_transformer_blocks.27': 'cpu', 'single_transformer_blocks.28': 'cpu', 'single_transformer_blocks.29': 'cpu', 'single_transformer_blocks.30': 'cpu', 'single_transformer_blocks.31': 'cpu', 'single_transformer_blocks.32': 'cpu', 'single_transformer_blocks.33': 'cpu', 'single_transformer_blocks.34': 'cpu', 'single_transformer_blocks.35': 'cpu', 'single_transformer_blocks.36': 'cpu', 'single_transformer_blocks.37': 'cpu', 'norm_out': 'cpu', 'proj_out': 'cpu'}

예를 들어, 아래 device_map은 single_transformer_blocks.10부터 single_transformer_blocks.20까지를 두 번째 GPU(1)에 배치합니다.

import torch
from diffusers import AutoModel

device_map = {
    'pos_embed': 0, 'time_text_embed': 0, 'context_embedder': 0, 'x_embedder': 0, 'transformer_blocks': 0, 'single_transformer_blocks.0': 0, 'single_transformer_blocks.1': 0, 'single_transformer_blocks.2': 0, 'single_transformer_blocks.3': 0, 'single_transformer_blocks.4': 0, 'single_transformer_blocks.5': 0, 'single_transformer_blocks.6': 0, 'single_transformer_blocks.7': 0, 'single_transformer_blocks.8': 0, 'single_transformer_blocks.9': 0, 'single_transformer_blocks.10': 1, 'single_transformer_blocks.11': 1, 'single_transformer_blocks.12': 1, 'single_transformer_blocks.13': 1, 'single_transformer_blocks.14': 1, 'single_transformer_blocks.15': 1, 'single_transformer_blocks.16': 1, 'single_transformer_blocks.17': 1, 'single_transformer_blocks.18': 1, 'single_transformer_blocks.19': 1, 'single_transformer_blocks.20': 1, 'single_transformer_blocks.21': 'cpu', 'single_transformer_blocks.22': 'cpu', 'single_transformer_blocks.23': 'cpu', 'single_transformer_blocks.24': 'cpu', 'single_transformer_blocks.25': 'cpu', 'single_transformer_blocks.26': 'cpu', 'single_transformer_blocks.27': 'cpu', 'single_transformer_blocks.28': 'cpu', 'single_transformer_blocks.29': 'cpu', 'single_transformer_blocks.30': 'cpu', 'single_transformer_blocks.31': 'cpu', 'single_transformer_blocks.32': 'cpu', 'single_transformer_blocks.33': 'cpu', 'single_transformer_blocks.34': 'cpu', 'single_transformer_blocks.35': 'cpu', 'single_transformer_blocks.36': 'cpu', 'single_transformer_blocks.37': 'cpu', 'norm_out': 'cpu', 'proj_out': 'cpu'
}

transformer = AutoModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev", 
    subfolder="transformer",
    device_map=device_map,
    dtype=torch.bfloat16
)

각 장치의 최대 메모리 사용량을 매핑하는 사전(dict)을 전달하여 한도를 강제합니다. max_memory에 장치가 없으면 무시되고, 파이프라인 구성 요소가 해당 장치에 분배되지 않습니다.

import torch
from diffusers import AutoModel, StableDiffusionXLPipeline

max_memory = {0:"1GB", 1:"1GB"}
pipeline = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    dtype=torch.float16,
    device_map="balanced",
    max_memory=max_memory
)

Diffusers는 기본적으로 모든 디바이스의 최대 메모리를 사용하지만, GPU에 맞지 않는 경우 아래 방법으로 단일 GPU를 사용하고 CPU로 오프로드해야 합니다.

enable_model_cpu_offload()은 단일 GPU에서만 작동하지만, 매우 큰 모델은 해당 GPU에 맞지 않을 수 있습니다
enable_sequential_cpu_offload()은 작동할 수 있지만 매우 느리며 단일 GPU로만 제한됩니다

device_map을 초기화하려면 reset_device_map() 메서드를 사용하세요. 디바이스 매핑이 적용된 파이프라인에서 .to(), enable_sequential_cpu_offload(), enable_model_cpu_offload()와 같은 메서드를 사용하려면 이 작업이 필요합니다.

pipeline.reset_device_map()

VAE 슬라이싱

VAE slicing은 큰 배치의 입력을 단일 배치의 데이터로 나누어 각각 별도로 처리함으로써 메모리를 절약합니다. 이 방법은 한 번에 여러 이미지를 생성할 때 가장 효과적입니다.

예를 들어, 이미지를 한 번에 4개 생성하는 경우 디코딩으로 인해 피크 활성화 메모리가 4배 증가합니다. VAE slicing은 4개 이미지를 한 번에 디코딩하는 대신 한 번에 1개씩만 디코딩하여 이를 줄입니다.

enable_slicing()을 호출하여 슬라이스된 VAE를 활성화합니다. 여러 이미지 배치의 디코딩 시 성능이 약간 향상되며, 단일 이미지 배치에는 성능에 영향을 미치지 않습니다.

import torch
from diffusers import AutoModel, StableDiffusionXLPipeline

pipeline = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    dtype=torch.float16,
).to("cuda")  # or "mps", "xpu", "cpu"
pipeline.vae.enable_slicing()
pipeline(["An astronaut riding a horse on Mars"]*32).images[0]
print(f"Max memory reserved: {torch.cuda.max_memory_allocated() / 1024**3:.2f} GB")

AutoencoderKLWan 및 AsymmetricAutoencoderKL 클래스는 슬라이싱을 지원하지 않습니다.

VAE 타일링

VAE tiling은 이미지를 한 번에 처리하는 대신 더 작은 겹치는 타일로 나누어 처리함으로써 메모리를 절약합니다. GPU가 한 번에 하나의 타일만 처리하기 때문에 피크 메모리 사용량도 줄어듭니다.

enable_tiling()을 호출하여 VAE 타일링을 활성화합니다. 타일이 개별적으로 디코딩되므로 생성된 이미지에는 타일 간에 약간의 톤 차이가 있을 수 있지만, 타일 사이에 눈에 띄는 이음새는 나타나지 않아야 합니다. 미리 지정된(그러나 설정 가능한) 한계보다 낮은 해상도에서는 타일링이 비활성화됩니다. 예를 들어, StableDiffusionPipeline의 VAE에서 이 한계는 512x512입니다.

import torch
from diffusers import AutoPipelineForImage2Image
from diffusers.utils import load_image

pipeline = AutoPipelineForImage2Image.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", dtype=torch.float16
).to("cuda")  # or "mps", "xpu", "cpu"
pipeline.vae.enable_tiling()

init_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/img2img-sdxl-init.png")
prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
pipeline(prompt, image=init_image, strength=0.5).images[0]
print(f"Max memory reserved: {torch.cuda.max_memory_allocated() / 1024**3:.2f} GB")

AutoencoderKLWan과 AsymmetricAutoencoderKL은 타일링을 지원하지 않습니다.

오프로딩

Offloading 전략은 현재 활성화되지 않은 레이어 또는 모델을 CPU로 이동시켜 GPU 메모리 증가를 방지합니다. 이러한 전략은 양자화(quantization) 및 torch.compile과 결합하여 추론 속도와 메모리 사용량 사이의 균형을 맞출 수 있습니다.

자세한 내용은 Compile and offloading quantized models 가이드를 참조하세요.

CPU 오프로딩

CPU offloading은 가중치를 GPU에서 CPU로 선택적으로 이동합니다. 특정 구성 요소가 필요할 때 GPU로 전송되고, 필요하지 않을 때 CPU로 이동합니다. 이 방법은 전체 모델이 아니라 서브모듈 단위로 작동합니다. GPU에 전체 모델을 저장하지 않도록 하여 메모리를 절약합니다.

CPU offloading은 메모리 사용량을 크게 줄여주지만, 서브모듈이 장치 간에 여러 번 오가야 하므로 매우 느립니다. 이러한 느림으로 인해 실제로 사용하기 어려울 수 있습니다.

enable_sequential_cpu_offload()를 호출하기 전에 파이프라인을 CUDA로 이동하지 마세요. 그렇지 않으면 절약되는 메모리 양이 매우 적습니다(자세한 내용은 이 이슈를 참조하세요). 이는 모델에 훅을 설치하는 상태 유지(stateful) 연산입니다.

파이프라인에서 이를 활성화하려면 enable_sequential_cpu_offload()를 호출하세요.

import torch
from diffusers import DiffusionPipeline

pipeline = DiffusionPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-schnell", dtype=torch.bfloat16
)
pipeline.enable_sequential_cpu_offload()

pipeline(
    prompt="An astronaut riding a horse on Mars",
    guidance_scale=0.,
    height=768,
    width=1360,
    num_inference_steps=4,
    max_sequence_length=256,
).images[0]
print(f"Max memory reserved: {torch.cuda.max_memory_allocated() / 1024**3:.2f} GB")

모델 오프로딩

Model offloading은 일부 레이어나 모델 구성 요소를 선택적으로 이동하는 대신 모델 전체를 GPU로 이동합니다. 파이프라인의 주요 모델 중 하나, 일반적으로 텍스트 인코더, UNet, VAE가 GPU에 배치되고 나머지 구성 요소는 CPU에 유지됩니다. 여러 번 실행되는 UNet과 같은 구성 요소는 완전히 종료되고 더 이상 필요하지 않을 때까지 GPU에 유지됩니다. 이를 통해 CPU offloading의 통신 오버헤드가 제거되며, model offloading은 더 빠른 대안으로 작용합니다. 단, 메모리 절감 효과가 크지 않다는 것이 트레이드오프입니다.

훅이 설치된 후 파이프라인 외부에서 모델이 재사용되는 경우(자세한 내용은 훅 제거를 참조)를 염두에 두세요. 이 경우 모델을 올바르게 오프로드하려면 파이프라인 전체와 모델을 예상된 순서대로 실행해야 합니다. 이는 모델에 훅을 설치하는 상태 유지(stateful) 연산입니다.

파이프라인에서 이를 활성화하려면 enable_model_cpu_offload()를 호출하세요.

import torch
from diffusers import DiffusionPipeline

pipeline = DiffusionPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-schnell", dtype=torch.bfloat16
)
pipeline.enable_model_cpu_offload()

pipeline(
    prompt="An astronaut riding a horse on Mars",
    guidance_scale=0.,
    height=768,
    width=1360,
    num_inference_steps=4,
    max_sequence_length=256,
).images[0]
print(f"Max memory reserved: {torch.cuda.max_memory_allocated() / 1024**3:.2f} GB")

enable_model_cpu_offload()는 encode_prompt() 메서드를 단독으로 사용해 텍스트 인코더의 hidden state를 생성할 때도 유용합니다.

그룹 오프로딩

Group offloading은 내부 레이어 그룹(torch.nn.ModuleList 또는 torch.nn.Sequential)을 CPU로 이동합니다. model offloading보다 메모리 사용량이 적고, 통신 오버헤드를 줄이기 때문에 CPU offloading보다 빠릅니다.

forward 구현에 입력의 가중치 기반 디바이스 캐스팅이 포함된 경우, 그룹 오프로딩의 디바이스 캐스팅 메커니즘과 충돌할 수 있으므로 모든 모델에서 그룹 오프로딩이 작동하지 않을 수 있습니다.

offload_type 매개변수를 block_level 또는 leaf_level로 설정하여 그룹 오프로딩을 활성화합니다.

block_level 오프로딩은 num_blocks_per_group 파라미터를 기반으로 레이어 그룹을 오프로드합니다. 예를 들어, 40개 레이어를 가진 모델에서 num_blocks_per_group=2인 경우, 한 번에 2개 레이어가 온로드 및 오프로드됩니다(총 20회 온로드/오프로드). 이를 통해 메모리 요구량이 크게 줄어듭니다.
leaf_level은 가장 낮은 수준에서 개별 레이어를 오프로드하며 CPU offloading과 동일합니다. 하지만 추론 속도를 포기하지 않으면서 스트림을 사용하면 더 빠르게 만들 수 있습니다.

Group offloading은 전체 파이프라인 또는 개별 모델에 적용할 수 있습니다. 전체 파이프라인에 group offloading을 적용하는 것이 가장 간단한 방법이며, 개별 모델에 선택적으로 적용하면 사용자가 모델별로 다른 offloading 기법을 사용할 수 있어 유연성이 높아집니다.

pipeline

모델

CUDA 스트림

CUDA 디바이스에서 비동기 데이터 전송 스트림을 지원하는 경우 use_stream 파라미터를 활성화하여 CPU 오프로딩 대비 전체 실행 시간을 줄일 수 있습니다. 레이어 프리페칭을 사용하여 데이터 전송과 계산을 겹치게 합니다. 현재 레이어가 실행되는 동안 다음으로 실행될 레이어를 GPU로 로드합니다. CPU 메모리가 크게 증가할 수 있으므로 모델 크기의 2배에 해당하는 메모리가 있는지 확인하세요.

메모리 사용량이 약간 증가하는 대신 더 큰 속도 향상을 얻으려면 record_stream=True를 설정하세요. 자세한 내용은 torch.Tensor.record_stream 문서를 참고하세요.

tiling이 활성화된 VAE에서 use_stream=True를 사용할 경우, device mismatch 오류를 방지하기 위해 추론 전에 dummy forward pass( dummy inputs를 사용해도 가능)를 수행해야 합니다. 모든 구현에서 이 방법이 작동하지 않을 수 있으므로, 문제가 발생하면 issue를 열어 주세요.

use_stream이 활성화된 상태에서 block_level 그룹 오프로딩을 사용하는 경우, num_blocks_per_group 매개변수는 1로 설정해야 하며, 그렇지 않으면 경고가 발생합니다.

pipeline.transformer.enable_group_offload(onload_device=onload_device, offload_device=offload_device, offload_type="leaf_level", use_stream=True, record_stream=True)

low_cpu_mem_usage 매개변수를 True로 설정하면 그룹 오프로딩 시 스트림을 사용하는 동안 CPU 메모리 사용량을 줄일 수 있습니다. leaf_level 오프로딩을 사용하거나 CPU 메모리가 병목 현상을 일으키는 경우에 가장 적합합니다. 미리 고정(pinning)하는 대신 필요할 때 고정된 텐서를 생성함으로써 메모리를 절약합니다. 그러나 이로 인해 전체 실행 시간이 증가할 수 있습니다.

디스크로 오프로딩

Group offloading은 모델 크기에 따라 상당한 시스템 메모리를 소모할 수 있습니다. 메모리가 제한된 시스템에서는 보조 메모리로 디스크에 group offloading을 적용해 보세요.

모델을 디스크로 오프로드하려면 enable_group_offload() 또는 apply_group_offloading()에서 offload_to_disk_path 인수를 설정하세요.

pipeline.transformer.enable_group_offload(onload_device=onload_device, offload_device=offload_device, offload_type="leaf_level", offload_to_disk_path="path/to/disk")

apply_group_offloading(pipeline.text_encoder, onload_device=onload_device, offload_type="block_level", num_blocks_per_group=2, offload_to_disk_path="path/to/disk")

속도와 메모리의 트레이드오프를 비교하려면 이 두 표를 참조하세요.

레이어별 캐스팅

레이어별 캐스팅을 그룹 오프로딩과 결합하여 메모리를 더 많이 절약할 수 있습니다.

Layerwise casting은 가중치를 더 작은 데이터 형식(예: torch.float8_e4m3fn 및 torch.float8_e5m2)으로 저장하여 메모리 사용량을 줄이고, 계산 시 torch.float16 또는 torch.bfloat16와 같은 더 높은 정밀도로 업캐스팅합니다. fp8로 저장할 경우 생성 품질이 저하될 수 있으므로 특정 레이어(정규화 및 변조 관련 가중치)는 건너뜁니다.

forward 구현에 가중치의 내부 타입 캐스팅이 포함된 경우, 모든 모델에서 레이어별 캐스팅이 작동하지 않을 수 있습니다. 현재 레이어별 캐스팅 구현은 forward pass가 가중치 정밀도와 독립적이며, 입력 데이터 타입은 항상 compute_dtype으로 지정된다고 가정합니다(호환되지 않는 구현은 여기를 참조).

PEFT 레이어가 포함된 커스텀 모델링 구현에서 레이어별 캐스팅이 실패할 수도 있습니다. 일부 체크가 제공되지만, 광범위하게 테스트되지 않았으며 모든 경우에 동작이 보장되지 않습니다.

저장 및 계산 데이터형을 설정하려면 enable_layerwise_casting()을 호출합니다.

import torch
from diffusers import CogVideoXPipeline, CogVideoXTransformer3DModel
from diffusers.utils import export_to_video

transformer = CogVideoXTransformer3DModel.from_pretrained(
    "THUDM/CogVideoX-5b",
    subfolder="transformer",
    dtype=torch.bfloat16
)
transformer.enable_layerwise_casting(storage_dtype=torch.float8_e4m3fn, compute_dtype=torch.bfloat16)

pipeline = CogVideoXPipeline.from_pretrained("THUDM/CogVideoX-5b",
    transformer=transformer,
    dtype=torch.bfloat16
).to("cuda")  # or "mps", "xpu", "cpu"
prompt = (
    "A panda, dressed in a small, red jacket and a tiny hat, sits on a wooden stool in a serene bamboo forest. "
    "The panda's fluffy paws strum a miniature acoustic guitar, producing soft, melodic tunes. Nearby, a few other "
    "pandas gather, watching curiously and some clapping in rhythm. Sunlight filters through the tall bamboo, "
    "casting a gentle glow on the scene. The panda's face is expressive, showing concentration and joy as it plays. "
    "The background includes a small, flowing stream and vibrant green foliage, enhancing the peaceful and magical "
    "atmosphere of this unique musical performance."
)
video = pipeline(prompt=prompt, guidance_scale=6, num_inference_steps=50).frames[0]
print(f"Max memory reserved: {torch.cuda.max_memory_allocated() / 1024**3:.2f} GB")
export_to_video(video, "output.mp4", fps=8)

더 많은 제어와 유연성이 필요할 경우 apply_layerwise_casting() 메서드를 사용할 수도 있습니다. 특정 내부 모듈에 대해 호출하여 모델 레이어에 부분적으로 적용할 수 있습니다. 정규화 및 변조 레이어와 같이 피할 모듈을 지정하려면 skip_modules_pattern 또는 skip_modules_classes 매개변수를 사용하세요.

import torch
from diffusers import CogVideoXTransformer3DModel
from diffusers.hooks import apply_layerwise_casting

transformer = CogVideoXTransformer3DModel.from_pretrained(
    "THUDM/CogVideoX-5b",
    subfolder="transformer",
    dtype=torch.bfloat16
)

# skip the normalization layer
apply_layerwise_casting(
    transformer,
    storage_dtype=torch.float8_e4m3fn,
    compute_dtype=torch.bfloat16,
    skip_modules_classes=["norm"],
    non_blocking=True,
)

torch.channels_last

torch.channels_last는 텐서의 저장 방식을 (batch size, channels, height, width)에서 (batch size, height, width, channels)로 변경합니다. 이를 통해 텐서가 메모리에 저장된 순서와 하드웨어가 텐서에 접근하는 방식이 일치하게 되어, 픽셀 값에 접근할 때 메모리에서 여기저기 건너뛰는 작업을 피할 수 있습니다.

현재 모든 연산자가 channels-last 형식을 지원하지는 않으며, 이로 인해 성능이 저하될 수 있지만, 여전히 시도해 볼 가치가 있습니다.

print(pipeline.unet.conv_out.state_dict()["weight"].stride())  # (2880, 9, 3, 1)
pipeline.unet.to(memory_format=torch.channels_last)  # in-place operation
print(
    pipeline.unet.conv_out.state_dict()["weight"].stride()
)  # (2880, 1, 960, 320) having a stride of 1 for the 2nd dimension proves that it works

메모리 효율적인 어텐션

Diffusers는 set_attention_backend()를 통해 여러 메모리 효율적인 어텐션 백엔드(FlashAttention, xFormers, SageAttention 등)를 지원합니다. 이들을 전환하는 방법은 Attention backends 가이드를 참고하세요.

Update on GitHub