From 068daa1e0915e01141a4b6699f4218f014e5b6c3 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Thu, 27 Mar 2025 09:13:10 -0400 Subject: [PATCH] add manual module move for video models Signed-off-by: Vladimir Mandic --- modules/video_models/video_utils.py | 3 ++- modules/video_models/video_vae.py | 1 + 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/modules/video_models/video_utils.py b/modules/video_models/video_utils.py index bb971d9d5..2b996a7e2 100644 --- a/modules/video_models/video_utils.py +++ b/modules/video_models/video_utils.py @@ -1,6 +1,6 @@ import os import time -from modules import shared, sd_models, timer, errors +from modules import shared, sd_models, timer, errors, devices debug = shared.log.trace if os.environ.get('SD_VIDEO_DEBUG', None) is not None else lambda *args, **kwargs: None @@ -31,6 +31,7 @@ def set_prompt(p): def hijack_encode_prompt(*args, **kwargs): t0 = time.time() try: + sd_models.move_model(shared.sd_model.text_encoder, devices.device) res = shared.sd_model.orig_encode_prompt(*args, **kwargs) except Exception as e: shared.log.error(f'Video encode: {e}') diff --git a/modules/video_models/video_vae.py b/modules/video_models/video_vae.py index 8adbc939e..5a107f616 100644 --- a/modules/video_models/video_vae.py +++ b/modules/video_models/video_vae.py @@ -57,6 +57,7 @@ def hijack_vae_decode(*args, **kwargs): if res is None: shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model, exclude=['vae']) try: + sd_models.move_model(shared.sd_model.vae, devices.device) if torch.is_tensor(args[0]): latent = args[0] latent = latent.to(device=devices.device, dtype=shared.sd_model.vae.dtype) # upcast to vae dtype