add offload mode to controlnets and framepack

Signed-off-by: Vladimir Mandic <mandic00@live.com>
This commit is contained in:
Vladimir Mandic
2025-10-22 15:40:52 -04:00
parent 3b89094ac1
commit 87fc6c7913
3 changed files with 29 additions and 6 deletions
+8
View File
@@ -172,6 +172,7 @@ class Processor():
return f' Processor(id={self.processor_id} model={self.model.__class__.__name__})' if self.processor_id and self.model else ''
def reset(self, processor_id: str = None):
from modules.shared import opts
if self.model is not None:
debug(f'Control Processor unloaded: id="{self.processor_id}"')
self.model = None
@@ -180,6 +181,13 @@ class Processor():
# self.override = None
# devices.torch_gc()
self.load_config = { 'cache_dir': cache_dir }
if opts.offline_mode:
self.load_config["local_files_only"] = True
os.environ['HF_HUB_OFFLINE'] = '1'
else:
os.environ.pop('HF_HUB_OFFLINE', None)
os.unsetenv('HF_HUB_OFFLINE')
def config(self, processor_id = None):
if processor_id is not None:
+6
View File
@@ -205,6 +205,12 @@ class ControlNet():
self.load_config = { 'cache_dir': cache_dir }
if load_config is not None:
self.load_config.update(load_config)
if opts.offline_mode:
self.load_config["local_files_only"] = True
os.environ['HF_HUB_OFFLINE'] = '1'
else:
os.environ.pop('HF_HUB_OFFLINE', None)
os.unsetenv('HF_HUB_OFFLINE')
if model_id is not None:
self.load()
+15 -6
View File
@@ -1,3 +1,4 @@
import os
import time
from modules import shared, devices, errors, sd_models, sd_checkpoint, model_quant
@@ -128,23 +129,31 @@ def load_model(variant:str=None, pipeline:str=None, text_encoder:str=None, text_
sd_models.hf_auth_check(model["text_encoder"]["repo"])
sd_models.hf_auth_check(model["text_encoder_2"]["repo"])
offline_config = {}
if shared.opts.offline_mode:
offline_config["local_files_only"] = True
os.environ['HF_HUB_OFFLINE'] = '1'
else:
os.environ.pop('HF_HUB_OFFLINE', None)
os.unsetenv('HF_HUB_OFFLINE')
shared.log.debug(f'FramePack load: module=llm {model["text_encoder"]}')
load_args, quant_args = model_quant.get_dit_args({}, module='TE', device_map=True)
text_encoder = LlamaModel.from_pretrained(model["text_encoder"]["repo"], subfolder=model["text_encoder"]["subfolder"], cache_dir=shared.opts.hfcache_dir, **load_args, **quant_args)
text_encoder = LlamaModel.from_pretrained(model["text_encoder"]["repo"], subfolder=model["text_encoder"]["subfolder"], cache_dir=shared.opts.hfcache_dir, **load_args, **quant_args, **offline_config)
tokenizer = LlamaTokenizerFast.from_pretrained(model["tokenizer"]["repo"], subfolder=model["tokenizer"]["subfolder"], cache_dir=shared.opts.hfcache_dir)
text_encoder.requires_grad_(False)
text_encoder.eval()
sd_models.move_model(text_encoder, devices.cpu)
shared.log.debug(f'FramePack load: module=te {model["text_encoder_2"]}')
text_encoder_2 = CLIPTextModel.from_pretrained(model["text_encoder_2"]["repo"], subfolder=model["text_encoder_2"]["subfolder"], torch_dtype=devices.dtype, cache_dir=shared.opts.hfcache_dir)
text_encoder_2 = CLIPTextModel.from_pretrained(model["text_encoder_2"]["repo"], subfolder=model["text_encoder_2"]["subfolder"], torch_dtype=devices.dtype, cache_dir=shared.opts.hfcache_dir, **offline_config)
tokenizer_2 = CLIPTokenizer.from_pretrained(model["pipeline"]["repo"], subfolder='tokenizer_2', cache_dir=shared.opts.hfcache_dir)
text_encoder_2.requires_grad_(False)
text_encoder_2.eval()
sd_models.move_model(text_encoder_2, devices.cpu)
shared.log.debug(f'FramePack load: module=vae {model["vae"]}')
vae = AutoencoderKLHunyuanVideo.from_pretrained(model["vae"]["repo"], subfolder=model["vae"]["subfolder"], torch_dtype=devices.dtype, cache_dir=shared.opts.hfcache_dir)
vae = AutoencoderKLHunyuanVideo.from_pretrained(model["vae"]["repo"], subfolder=model["vae"]["subfolder"], torch_dtype=devices.dtype, cache_dir=shared.opts.hfcache_dir, **offline_config)
vae.requires_grad_(False)
vae.eval()
vae.enable_slicing()
@@ -152,8 +161,8 @@ def load_model(variant:str=None, pipeline:str=None, text_encoder:str=None, text_
sd_models.move_model(vae, devices.cpu)
shared.log.debug(f'FramePack load: module=encoder {model["feature_extractor"]} model={model["image_encoder"]}')
feature_extractor = SiglipImageProcessor.from_pretrained(model["feature_extractor"]["repo"], subfolder=model["feature_extractor"]["subfolder"], cache_dir=shared.opts.hfcache_dir)
image_encoder = SiglipVisionModel.from_pretrained(model["image_encoder"]["repo"], subfolder=model["image_encoder"]["subfolder"], torch_dtype=devices.dtype, cache_dir=shared.opts.hfcache_dir)
feature_extractor = SiglipImageProcessor.from_pretrained(model["feature_extractor"]["repo"], subfolder=model["feature_extractor"]["subfolder"], cache_dir=shared.opts.hfcache_dir, **offline_config)
image_encoder = SiglipVisionModel.from_pretrained(model["image_encoder"]["repo"], subfolder=model["image_encoder"]["subfolder"], torch_dtype=devices.dtype, cache_dir=shared.opts.hfcache_dir, **offline_config)
image_encoder.requires_grad_(False)
image_encoder.eval()
sd_models.move_model(image_encoder, devices.cpu)
@@ -161,7 +170,7 @@ def load_model(variant:str=None, pipeline:str=None, text_encoder:str=None, text_
shared.log.debug(f'FramePack load: module=transformer {model["transformer"]}')
dit_repo = model["transformer"]["repo"]
load_args, quant_args = model_quant.get_dit_args({}, module='Model', device_map=True)
transformer = HunyuanVideoTransformer3DModelPacked.from_pretrained(dit_repo, subfolder=model["transformer"]["subfolder"], cache_dir=shared.opts.hfcache_dir, **load_args, **quant_args)
transformer = HunyuanVideoTransformer3DModelPacked.from_pretrained(dit_repo, subfolder=model["transformer"]["subfolder"], cache_dir=shared.opts.hfcache_dir, **load_args, **quant_args, **offline_config)
transformer.high_quality_fp32_output_for_inference = False
transformer.requires_grad_(False)
transformer.eval()