File size: 2,556 Bytes
fa2d87b 478202c fa2d87b 478202c fa2d87b 478202c fa2d87b 478202c fa2d87b 478202c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 | #!/usr/bin/env python3
from __future__ import annotations
from collections.abc import Callable
from accelerate import cpu_offload
class _ExecutionDeviceHint:
execution_device = "cuda"
offload = False
hooks: tuple = ()
def move_registered_buffers(model, device: str) -> int:
"""Repair registered buffers missed by quantized-model placement wrappers."""
moved = 0
for module in model.modules():
for name, buffer in getattr(module, "_buffers", {}).items():
if buffer is not None and str(buffer.device) != device:
module._buffers[name] = buffer.to(device)
moved += 1
return moved
def install_h3_rope_runtime_alignment(transformer) -> None:
rope = transformer.rope
original_forward = rope.forward
def aligned_forward(position_ids):
inv_freq = rope.inv_freq
if str(inv_freq.device) != str(position_ids.device):
rope._buffers["inv_freq"] = inv_freq.to(position_ids.device)
return original_forward(position_ids)
rope.forward = aligned_forward
def install_manual_h3_stage_offload(
encoder_step_cls,
*,
text_encoder,
transformer,
empty_cuda_cache: Callable[[], None],
place_transformer: bool = True,
sequential_text_encoder: bool = False,
) -> None:
"""Run conditioning on CUDA, then free it before placing the denoiser."""
install_h3_rope_runtime_alignment(transformer)
original_encode_prompt = encoder_step_cls.encode_prompt
original_call = encoder_step_cls.__call__
def encode_prompt_on_cuda(
components,
prompt,
images=None,
device=None,
dtype=None,
):
return original_encode_prompt(
components,
prompt,
images,
device="cuda",
dtype=dtype,
)
def call_then_place_denoiser(step, components, state):
result = original_call(step, components, state)
if not sequential_text_encoder:
text_encoder.to("cpu")
text_encoder._hf_hook = _ExecutionDeviceHint()
empty_cuda_cache()
if place_transformer:
transformer.to("cuda")
move_registered_buffers(transformer, "cuda")
return result
encoder_step_cls.encode_prompt = staticmethod(encode_prompt_on_cuda)
encoder_step_cls.__call__ = call_then_place_denoiser
if sequential_text_encoder:
cpu_offload(text_encoder, execution_device="cuda", offload_buffers=True)
else:
text_encoder.to("cuda")
|