fred-dev/comfy_ui_ali
0
1from __future__ import annotations2import json3import torch4from enum import Enum5import logging6 7from comfy import model_management8from comfy.utils import ProgressBar9from .ldm.models.autoencoder import AutoencoderKL, AutoencodingEngine10from .ldm.cascade.stage_a import StageA11from .ldm.cascade.stage_c_coder import StageC_coder12from .ldm.audio.autoencoder import AudioOobleckVAE13import comfy.ldm.genmo.vae.model14import comfy.ldm.lightricks.vae.causal_video_autoencoder15import comfy.ldm.cosmos.vae16import comfy.ldm.wan.vae17import yaml18import math19 20import comfy.utils21 22from . import clip_vision23from . import gligen24from . import diffusers_convert25from . import model_detection26 27from . import sd1_clip28from . import sdxl_clip29import comfy.text_encoders.sd2_clip30import comfy.text_encoders.sd3_clip31import comfy.text_encoders.sa_t532import comfy.text_encoders.aura_t533import comfy.text_encoders.pixart_t534import comfy.text_encoders.hydit35import comfy.text_encoders.flux36import comfy.text_encoders.long_clipl37import comfy.text_encoders.genmo38import comfy.text_encoders.lt39import comfy.text_encoders.hunyuan_video40import comfy.text_encoders.cosmos41import comfy.text_encoders.lumina242import comfy.text_encoders.wan43 44import comfy.model_patcher45import comfy.lora46import comfy.lora_convert47import comfy.hooks48import comfy.t2i_adapter.adapter49import comfy.taesd.taesd50 51import comfy.ldm.flux.redux52 53def load_lora_for_models(model, clip, lora, strength_model, strength_clip):54 key_map = {}55 if model is not None:56 key_map = comfy.lora.model_lora_keys_unet(model.model, key_map)57 if clip is not None:58 key_map = comfy.lora.model_lora_keys_clip(clip.cond_stage_model, key_map)59 60 lora = comfy.lora_convert.convert_lora(lora)61 loaded = comfy.lora.load_lora(lora, key_map)62 if model is not None:63 new_modelpatcher = model.clone()64 k = new_modelpatcher.add_patches(loaded, strength_model)65 else:66 k = ()67 new_modelpatcher = None68 69 if clip is not None:70 new_clip = clip.clone()71 k1 = new_clip.add_patches(loaded, strength_clip)72 else:73 k1 = ()74 new_clip = None75 k = set(k)76 k1 = set(k1)77 for x in loaded:78 if (x not in k) and (x not in k1):79 logging.warning("NOT LOADED {}".format(x))80 81 return (new_modelpatcher, new_clip)82 83 84class CLIP:85 def __init__(self, target=None, embedding_directory=None, no_init=False, tokenizer_data={}, parameters=0, model_options={}):86 if no_init:87 return88 params = target.params.copy()89 clip = target.clip90 tokenizer = target.tokenizer91 92 load_device = model_options.get("load_device", model_management.text_encoder_device())93 offload_device = model_options.get("offload_device", model_management.text_encoder_offload_device())94 dtype = model_options.get("dtype", None)95 if dtype is None:96 dtype = model_management.text_encoder_dtype(load_device)97 98 params['dtype'] = dtype99 params['device'] = model_options.get("initial_device", model_management.text_encoder_initial_device(load_device, offload_device, parameters * model_management.dtype_size(dtype)))100 params['model_options'] = model_options101 102 self.cond_stage_model = clip(**(params))103 104 for dt in self.cond_stage_model.dtypes:105 if not model_management.supports_cast(load_device, dt):106 load_device = offload_device107 if params['device'] != offload_device:108 self.cond_stage_model.to(offload_device)109 logging.warning("Had to shift TE back.")110 111 self.tokenizer = tokenizer(embedding_directory=embedding_directory, tokenizer_data=tokenizer_data)112 self.patcher = comfy.model_patcher.ModelPatcher(self.cond_stage_model, load_device=load_device, offload_device=offload_device)113 self.patcher.hook_mode = comfy.hooks.EnumHookMode.MinVram114 self.patcher.is_clip = True115 self.apply_hooks_to_conds = None116 if params['device'] == load_device:117 model_management.load_models_gpu([self.patcher], force_full_load=True)118 self.layer_idx = None119 self.use_clip_schedule = False120 logging.info("CLIP/text encoder model load device: {}, offload device: {}, current: {}, dtype: {}".format(load_device, offload_device, params['device'], dtype))121 122 def clone(self):123 n = CLIP(no_init=True)124 n.patcher = self.patcher.clone()125 n.cond_stage_model = self.cond_stage_model126 n.tokenizer = self.tokenizer127 n.layer_idx = self.layer_idx128 n.use_clip_schedule = self.use_clip_schedule129 n.apply_hooks_to_conds = self.apply_hooks_to_conds130 return n131 132 def add_patches(self, patches, strength_patch=1.0, strength_model=1.0):133 return self.patcher.add_patches(patches, strength_patch, strength_model)134 135 def clip_layer(self, layer_idx):136 self.layer_idx = layer_idx137 138 def tokenize(self, text, return_word_ids=False, **kwargs):139 return self.tokenizer.tokenize_with_weights(text, return_word_ids, **kwargs)140 141 def add_hooks_to_dict(self, pooled_dict: dict[str]):142 if self.apply_hooks_to_conds:143 pooled_dict["hooks"] = self.apply_hooks_to_conds144 return pooled_dict145 146 def encode_from_tokens_scheduled(self, tokens, unprojected=False, add_dict: dict[str]={}, show_pbar=True):147 all_cond_pooled: list[tuple[torch.Tensor, dict[str]]] = []148 all_hooks = self.patcher.forced_hooks149 if all_hooks is None or not self.use_clip_schedule:150 # if no hooks or shouldn't use clip schedule, do unscheduled encode_from_tokens and perform add_dict151 return_pooled = "unprojected" if unprojected else True152 pooled_dict = self.encode_from_tokens(tokens, return_pooled=return_pooled, return_dict=True)153 cond = pooled_dict.pop("cond")154 # add/update any keys with the provided add_dict155 pooled_dict.update(add_dict)156 all_cond_pooled.append([cond, pooled_dict])157 else:158 scheduled_keyframes = all_hooks.get_hooks_for_clip_schedule()159 160 self.cond_stage_model.reset_clip_options()161 if self.layer_idx is not None:162 self.cond_stage_model.set_clip_options({"layer": self.layer_idx})163 if unprojected:164 self.cond_stage_model.set_clip_options({"projected_pooled": False})165 166 self.load_model()167 all_hooks.reset()168 self.patcher.patch_hooks(None)169 if show_pbar:170 pbar = ProgressBar(len(scheduled_keyframes))171 172 for scheduled_opts in scheduled_keyframes:173 t_range = scheduled_opts[0]174 # don't bother encoding any conds outside of start_percent and end_percent bounds175 if "start_percent" in add_dict:176 if t_range[1] < add_dict["start_percent"]:177 continue178 if "end_percent" in add_dict:179 if t_range[0] > add_dict["end_percent"]:180 continue181 hooks_keyframes = scheduled_opts[1]182 for hook, keyframe in hooks_keyframes:183 hook.hook_keyframe._current_keyframe = keyframe184 # apply appropriate hooks with values that match new hook_keyframe185 self.patcher.patch_hooks(all_hooks)186 # perform encoding as normal187 o = self.cond_stage_model.encode_token_weights(tokens)188 cond, pooled = o[:2]189 pooled_dict = {"pooled_output": pooled}190 # add clip_start_percent and clip_end_percent in pooled191 pooled_dict["clip_start_percent"] = t_range[0]192 pooled_dict["clip_end_percent"] = t_range[1]193 # add/update any keys with the provided add_dict194 pooled_dict.update(add_dict)195 # add hooks stored on clip196 self.add_hooks_to_dict(pooled_dict)197 all_cond_pooled.append([cond, pooled_dict])198 if show_pbar:199 pbar.update(1)200 model_management.throw_exception_if_processing_interrupted()201 all_hooks.reset()202 return all_cond_pooled203 204 def encode_from_tokens(self, tokens, return_pooled=False, return_dict=False):205 self.cond_stage_model.reset_clip_options()206 207 if self.layer_idx is not None:208 self.cond_stage_model.set_clip_options({"layer": self.layer_idx})209 210 if return_pooled == "unprojected":211 self.cond_stage_model.set_clip_options({"projected_pooled": False})212 213 self.load_model()214 o = self.cond_stage_model.encode_token_weights(tokens)215 cond, pooled = o[:2]216 if return_dict:217 out = {"cond": cond, "pooled_output": pooled}218 if len(o) > 2:219 for k in o[2]:220 out[k] = o[2][k]221 self.add_hooks_to_dict(out)222 return out223 224 if return_pooled:225 return cond, pooled226 return cond227 228 def encode(self, text):229 tokens = self.tokenize(text)230 return self.encode_from_tokens(tokens)231 232 def load_sd(self, sd, full_model=False):233 if full_model:234 return self.cond_stage_model.load_state_dict(sd, strict=False)235 else:236 return self.cond_stage_model.load_sd(sd)237 238 def get_sd(self):239 sd_clip = self.cond_stage_model.state_dict()240 sd_tokenizer = self.tokenizer.state_dict()241 for k in sd_tokenizer:242 sd_clip[k] = sd_tokenizer[k]243 return sd_clip244 245 def load_model(self):246 model_management.load_model_gpu(self.patcher)247 return self.patcher248 249 def get_key_patches(self):250 return self.patcher.get_key_patches()251 252class VAE:253 def __init__(self, sd=None, device=None, config=None, dtype=None, metadata=None):254 if 'decoder.up_blocks.0.resnets.0.norm1.weight' in sd.keys(): #diffusers format255 sd = diffusers_convert.convert_vae_state_dict(sd)256 257 self.memory_used_encode = lambda shape, dtype: (1767 * shape[2] * shape[3]) * model_management.dtype_size(dtype) #These are for AutoencoderKL and need tweaking (should be lower)258 self.memory_used_decode = lambda shape, dtype: (2178 * shape[2] * shape[3] * 64) * model_management.dtype_size(dtype)259 self.downscale_ratio = 8260 self.upscale_ratio = 8261 self.latent_channels = 4262 self.latent_dim = 2263 self.output_channels = 3264 self.process_input = lambda image: image * 2.0 - 1.0265 self.process_output = lambda image: torch.clamp((image + 1.0) / 2.0, min=0.0, max=1.0)266 self.working_dtypes = [torch.bfloat16, torch.float32]267 268 self.downscale_index_formula = None269 self.upscale_index_formula = None270 271 if config is None:272 if "decoder.mid.block_1.mix_factor" in sd:273 encoder_config = {'double_z': True, 'z_channels': 4, 'resolution': 256, 'in_channels': 3, 'out_ch': 3, 'ch': 128, 'ch_mult': [1, 2, 4, 4], 'num_res_blocks': 2, 'attn_resolutions': [], 'dropout': 0.0}274 decoder_config = encoder_config.copy()275 decoder_config["video_kernel_size"] = [3, 1, 1]276 decoder_config["alpha"] = 0.0277 self.first_stage_model = AutoencodingEngine(regularizer_config={'target': "comfy.ldm.models.autoencoder.DiagonalGaussianRegularizer"},278 encoder_config={'target': "comfy.ldm.modules.diffusionmodules.model.Encoder", 'params': encoder_config},279 decoder_config={'target': "comfy.ldm.modules.temporal_ae.VideoDecoder", 'params': decoder_config})280 elif "taesd_decoder.1.weight" in sd:281 self.latent_channels = sd["taesd_decoder.1.weight"].shape[1]282 self.first_stage_model = comfy.taesd.taesd.TAESD(latent_channels=self.latent_channels)283 elif "vquantizer.codebook.weight" in sd: #VQGan: stage a of stable cascade284 self.first_stage_model = StageA()285 self.downscale_ratio = 4286 self.upscale_ratio = 4287 #TODO288 #self.memory_used_encode289 #self.memory_used_decode290 self.process_input = lambda image: image291 self.process_output = lambda image: image292 elif "backbone.1.0.block.0.1.num_batches_tracked" in sd: #effnet: encoder for stage c latent of stable cascade293 self.first_stage_model = StageC_coder()294 self.downscale_ratio = 32295 self.latent_channels = 16296 new_sd = {}297 for k in sd:298 new_sd["encoder.{}".format(k)] = sd[k]299 sd = new_sd300 elif "blocks.11.num_batches_tracked" in sd: #previewer: decoder for stage c latent of stable cascade301 self.first_stage_model = StageC_coder()302 self.latent_channels = 16303 new_sd = {}304 for k in sd:305 new_sd["previewer.{}".format(k)] = sd[k]306 sd = new_sd307 elif "encoder.backbone.1.0.block.0.1.num_batches_tracked" in sd: #combined effnet and previewer for stable cascade308 self.first_stage_model = StageC_coder()309 self.downscale_ratio = 32310 self.latent_channels = 16311 elif "decoder.conv_in.weight" in sd:312 #default SD1.x/SD2.x VAE parameters313 ddconfig = {'double_z': True, 'z_channels': 4, 'resolution': 256, 'in_channels': 3, 'out_ch': 3, 'ch': 128, 'ch_mult': [1, 2, 4, 4], 'num_res_blocks': 2, 'attn_resolutions': [], 'dropout': 0.0}314 315 if 'encoder.down.2.downsample.conv.weight' not in sd and 'decoder.up.3.upsample.conv.weight' not in sd: #Stable diffusion x4 upscaler VAE316 ddconfig['ch_mult'] = [1, 2, 4]317 self.downscale_ratio = 4318 self.upscale_ratio = 4319 320 self.latent_channels = ddconfig['z_channels'] = sd["decoder.conv_in.weight"].shape[1]321 if 'post_quant_conv.weight' in sd:322 self.first_stage_model = AutoencoderKL(ddconfig=ddconfig, embed_dim=sd['post_quant_conv.weight'].shape[1])323 else:324 self.first_stage_model = AutoencodingEngine(regularizer_config={'target': "comfy.ldm.models.autoencoder.DiagonalGaussianRegularizer"},325 encoder_config={'target': "comfy.ldm.modules.diffusionmodules.model.Encoder", 'params': ddconfig},326 decoder_config={'target': "comfy.ldm.modules.diffusionmodules.model.Decoder", 'params': ddconfig})327 elif "decoder.layers.1.layers.0.beta" in sd:328 self.first_stage_model = AudioOobleckVAE()329 self.memory_used_encode = lambda shape, dtype: (1000 * shape[2]) * model_management.dtype_size(dtype)330 self.memory_used_decode = lambda shape, dtype: (1000 * shape[2] * 2048) * model_management.dtype_size(dtype)331 self.latent_channels = 64332 self.output_channels = 2333 self.upscale_ratio = 2048334 self.downscale_ratio = 2048335 self.latent_dim = 1336 self.process_output = lambda audio: audio337 self.process_input = lambda audio: audio338 self.working_dtypes = [torch.float16, torch.bfloat16, torch.float32]339 elif "blocks.2.blocks.3.stack.5.weight" in sd or "decoder.blocks.2.blocks.3.stack.5.weight" in sd or "layers.4.layers.1.attn_block.attn.qkv.weight" in sd or "encoder.layers.4.layers.1.attn_block.attn.qkv.weight" in sd: #genmo mochi vae340 if "blocks.2.blocks.3.stack.5.weight" in sd:341 sd = comfy.utils.state_dict_prefix_replace(sd, {"": "decoder."})342 if "layers.4.layers.1.attn_block.attn.qkv.weight" in sd:343 sd = comfy.utils.state_dict_prefix_replace(sd, {"": "encoder."})344 self.first_stage_model = comfy.ldm.genmo.vae.model.VideoVAE()345 self.latent_channels = 12346 self.latent_dim = 3347 self.memory_used_decode = lambda shape, dtype: (1000 * shape[2] * shape[3] * shape[4] * (6 * 8 * 8)) * model_management.dtype_size(dtype)348 self.memory_used_encode = lambda shape, dtype: (1.5 * max(shape[2], 7) * shape[3] * shape[4] * (6 * 8 * 8)) * model_management.dtype_size(dtype)349 self.upscale_ratio = (lambda a: max(0, a * 6 - 5), 8, 8)350 self.upscale_index_formula = (6, 8, 8)351 self.downscale_ratio = (lambda a: max(0, math.floor((a + 5) / 6)), 8, 8)352 self.downscale_index_formula = (6, 8, 8)353 self.working_dtypes = [torch.float16, torch.float32]354 elif "decoder.up_blocks.0.res_blocks.0.conv1.conv.weight" in sd: #lightricks ltxv355 tensor_conv1 = sd["decoder.up_blocks.0.res_blocks.0.conv1.conv.weight"]356 version = 0357 if tensor_conv1.shape[0] == 512:358 version = 0359 elif tensor_conv1.shape[0] == 1024:360 version = 1361 if "encoder.down_blocks.1.conv.conv.bias" in sd:362 version = 2363 vae_config = None364 if metadata is not None and "config" in metadata:365 vae_config = json.loads(metadata["config"]).get("vae", None)366 self.first_stage_model = comfy.ldm.lightricks.vae.causal_video_autoencoder.VideoVAE(version=version, config=vae_config)367 self.latent_channels = 128368 self.latent_dim = 3369 self.memory_used_decode = lambda shape, dtype: (900 * shape[2] * shape[3] * shape[4] * (8 * 8 * 8)) * model_management.dtype_size(dtype)370 self.memory_used_encode = lambda shape, dtype: (70 * max(shape[2], 7) * shape[3] * shape[4]) * model_management.dtype_size(dtype)371 self.upscale_ratio = (lambda a: max(0, a * 8 - 7), 32, 32)372 self.upscale_index_formula = (8, 32, 32)373 self.downscale_ratio = (lambda a: max(0, math.floor((a + 7) / 8)), 32, 32)374 self.downscale_index_formula = (8, 32, 32)375 self.working_dtypes = [torch.bfloat16, torch.float32]376 elif "decoder.conv_in.conv.weight" in sd:377 ddconfig = {'double_z': True, 'z_channels': 4, 'resolution': 256, 'in_channels': 3, 'out_ch': 3, 'ch': 128, 'ch_mult': [1, 2, 4, 4], 'num_res_blocks': 2, 'attn_resolutions': [], 'dropout': 0.0}378 ddconfig["conv3d"] = True379 ddconfig["time_compress"] = 4380 self.upscale_ratio = (lambda a: max(0, a * 4 - 3), 8, 8)381 self.upscale_index_formula = (4, 8, 8)382 self.downscale_ratio = (lambda a: max(0, math.floor((a + 3) / 4)), 8, 8)383 self.downscale_index_formula = (4, 8, 8)384 self.latent_dim = 3385 self.latent_channels = ddconfig['z_channels'] = sd["decoder.conv_in.conv.weight"].shape[1]386 self.first_stage_model = AutoencoderKL(ddconfig=ddconfig, embed_dim=sd['post_quant_conv.weight'].shape[1])387 self.memory_used_decode = lambda shape, dtype: (1500 * shape[2] * shape[3] * shape[4] * (4 * 8 * 8)) * model_management.dtype_size(dtype)388 self.memory_used_encode = lambda shape, dtype: (900 * max(shape[2], 2) * shape[3] * shape[4]) * model_management.dtype_size(dtype)389 self.working_dtypes = [torch.bfloat16, torch.float16, torch.float32]390 elif "decoder.unpatcher3d.wavelets" in sd:391 self.upscale_ratio = (lambda a: max(0, a * 8 - 7), 8, 8)392 self.upscale_index_formula = (8, 8, 8)393 self.downscale_ratio = (lambda a: max(0, math.floor((a + 7) / 8)), 8, 8)394 self.downscale_index_formula = (8, 8, 8)395 self.latent_dim = 3396 self.latent_channels = 16397 ddconfig = {'z_channels': 16, 'latent_channels': self.latent_channels, 'z_factor': 1, 'resolution': 1024, 'in_channels': 3, 'out_channels': 3, 'channels': 128, 'channels_mult': [2, 4, 4], 'num_res_blocks': 2, 'attn_resolutions': [32], 'dropout': 0.0, 'patch_size': 4, 'num_groups': 1, 'temporal_compression': 8, 'spacial_compression': 8}398 self.first_stage_model = comfy.ldm.cosmos.vae.CausalContinuousVideoTokenizer(**ddconfig)399 #TODO: these values are a bit off because this is not a standard VAE400 self.memory_used_decode = lambda shape, dtype: (50 * shape[2] * shape[3] * shape[4] * (8 * 8 * 8)) * model_management.dtype_size(dtype)401 self.memory_used_encode = lambda shape, dtype: (50 * (round((shape[2] + 7) / 8) * 8) * shape[3] * shape[4]) * model_management.dtype_size(dtype)402 self.working_dtypes = [torch.bfloat16, torch.float32]403 elif "decoder.middle.0.residual.0.gamma" in sd:404 self.upscale_ratio = (lambda a: max(0, a * 4 - 3), 8, 8)405 self.upscale_index_formula = (4, 8, 8)406 self.downscale_ratio = (lambda a: max(0, math.floor((a + 3) / 4)), 8, 8)407 self.downscale_index_formula = (4, 8, 8)408 self.latent_dim = 3409 self.latent_channels = 16410 ddconfig = {"dim": 96, "z_dim": self.latent_channels, "dim_mult": [1, 2, 4, 4], "num_res_blocks": 2, "attn_scales": [], "temperal_downsample": [False, True, True], "dropout": 0.0}411 self.first_stage_model = comfy.ldm.wan.vae.WanVAE(**ddconfig)412 self.working_dtypes = [torch.bfloat16, torch.float16, torch.float32]413 self.memory_used_encode = lambda shape, dtype: 6000 * shape[3] * shape[4] * model_management.dtype_size(dtype)414 self.memory_used_decode = lambda shape, dtype: 7000 * shape[3] * shape[4] * (8 * 8) * model_management.dtype_size(dtype)415 else:416 logging.warning("WARNING: No VAE weights detected, VAE not initalized.")417 self.first_stage_model = None418 return419 else:420 self.first_stage_model = AutoencoderKL(**(config['params']))421 self.first_stage_model = self.first_stage_model.eval()422 423 m, u = self.first_stage_model.load_state_dict(sd, strict=False)424 if len(m) > 0:425 logging.warning("Missing VAE keys {}".format(m))426 427 if len(u) > 0:428 logging.debug("Leftover VAE keys {}".format(u))429 430 if device is None:431 device = model_management.vae_device()432 self.device = device433 offload_device = model_management.vae_offload_device()434 if dtype is None:435 dtype = model_management.vae_dtype(self.device, self.working_dtypes)436 self.vae_dtype = dtype437 self.first_stage_model.to(self.vae_dtype)438 self.output_device = model_management.intermediate_device()439 440 self.patcher = comfy.model_patcher.ModelPatcher(self.first_stage_model, load_device=self.device, offload_device=offload_device)441 logging.info("VAE load device: {}, offload device: {}, dtype: {}".format(self.device, offload_device, self.vae_dtype))442 443 def throw_exception_if_invalid(self):444 if self.first_stage_model is None:445 raise RuntimeError("ERROR: VAE is invalid: None\n\nIf the VAE is from a checkpoint loader node your checkpoint does not contain a valid VAE.")446 447 def vae_encode_crop_pixels(self, pixels):448 downscale_ratio = self.spacial_compression_encode()449 450 dims = pixels.shape[1:-1]451 for d in range(len(dims)):452 x = (dims[d] // downscale_ratio) * downscale_ratio453 x_offset = (dims[d] % downscale_ratio) // 2454 if x != dims[d]:455 pixels = pixels.narrow(d + 1, x_offset, x)456 return pixels457 458 def decode_tiled_(self, samples, tile_x=64, tile_y=64, overlap = 16):459 steps = samples.shape[0] * comfy.utils.get_tiled_scale_steps(samples.shape[3], samples.shape[2], tile_x, tile_y, overlap)460 steps += samples.shape[0] * comfy.utils.get_tiled_scale_steps(samples.shape[3], samples.shape[2], tile_x // 2, tile_y * 2, overlap)461 steps += samples.shape[0] * comfy.utils.get_tiled_scale_steps(samples.shape[3], samples.shape[2], tile_x * 2, tile_y // 2, overlap)462 pbar = comfy.utils.ProgressBar(steps)463 464 decode_fn = lambda a: self.first_stage_model.decode(a.to(self.vae_dtype).to(self.device)).float()465 output = self.process_output(466 (comfy.utils.tiled_scale(samples, decode_fn, tile_x // 2, tile_y * 2, overlap, upscale_amount = self.upscale_ratio, output_device=self.output_device, pbar = pbar) +467 comfy.utils.tiled_scale(samples, decode_fn, tile_x * 2, tile_y // 2, overlap, upscale_amount = self.upscale_ratio, output_device=self.output_device, pbar = pbar) +468 comfy.utils.tiled_scale(samples, decode_fn, tile_x, tile_y, overlap, upscale_amount = self.upscale_ratio, output_device=self.output_device, pbar = pbar))469 / 3.0)470 return output471 472 def decode_tiled_1d(self, samples, tile_x=128, overlap=32):473 decode_fn = lambda a: self.first_stage_model.decode(a.to(self.vae_dtype).to(self.device)).float()474 return self.process_output(comfy.utils.tiled_scale_multidim(samples, decode_fn, tile=(tile_x,), overlap=overlap, upscale_amount=self.upscale_ratio, out_channels=self.output_channels, output_device=self.output_device))475 476 def decode_tiled_3d(self, samples, tile_t=999, tile_x=32, tile_y=32, overlap=(1, 8, 8)):477 decode_fn = lambda a: self.first_stage_model.decode(a.to(self.vae_dtype).to(self.device)).float()478 return self.process_output(comfy.utils.tiled_scale_multidim(samples, decode_fn, tile=(tile_t, tile_x, tile_y), overlap=overlap, upscale_amount=self.upscale_ratio, out_channels=self.output_channels, index_formulas=self.upscale_index_formula, output_device=self.output_device))479 480 def encode_tiled_(self, pixel_samples, tile_x=512, tile_y=512, overlap = 64):481 steps = pixel_samples.shape[0] * comfy.utils.get_tiled_scale_steps(pixel_samples.shape[3], pixel_samples.shape[2], tile_x, tile_y, overlap)482 steps += pixel_samples.shape[0] * comfy.utils.get_tiled_scale_steps(pixel_samples.shape[3], pixel_samples.shape[2], tile_x // 2, tile_y * 2, overlap)483 steps += pixel_samples.shape[0] * comfy.utils.get_tiled_scale_steps(pixel_samples.shape[3], pixel_samples.shape[2], tile_x * 2, tile_y // 2, overlap)484 pbar = comfy.utils.ProgressBar(steps)485 486 encode_fn = lambda a: self.first_stage_model.encode((self.process_input(a)).to(self.vae_dtype).to(self.device)).float()487 samples = comfy.utils.tiled_scale(pixel_samples, encode_fn, tile_x, tile_y, overlap, upscale_amount = (1/self.downscale_ratio), out_channels=self.latent_channels, output_device=self.output_device, pbar=pbar)488 samples += comfy.utils.tiled_scale(pixel_samples, encode_fn, tile_x * 2, tile_y // 2, overlap, upscale_amount = (1/self.downscale_ratio), out_channels=self.latent_channels, output_device=self.output_device, pbar=pbar)489 samples += comfy.utils.tiled_scale(pixel_samples, encode_fn, tile_x // 2, tile_y * 2, overlap, upscale_amount = (1/self.downscale_ratio), out_channels=self.latent_channels, output_device=self.output_device, pbar=pbar)490 samples /= 3.0491 return samples492 493 def encode_tiled_1d(self, samples, tile_x=128 * 2048, overlap=32 * 2048):494 encode_fn = lambda a: self.first_stage_model.encode((self.process_input(a)).to(self.vae_dtype).to(self.device)).float()495 return comfy.utils.tiled_scale_multidim(samples, encode_fn, tile=(tile_x,), overlap=overlap, upscale_amount=(1/self.downscale_ratio), out_channels=self.latent_channels, output_device=self.output_device)496 497 def encode_tiled_3d(self, samples, tile_t=9999, tile_x=512, tile_y=512, overlap=(1, 64, 64)):498 encode_fn = lambda a: self.first_stage_model.encode((self.process_input(a)).to(self.vae_dtype).to(self.device)).float()499 return comfy.utils.tiled_scale_multidim(samples, encode_fn, tile=(tile_t, tile_x, tile_y), overlap=overlap, upscale_amount=self.downscale_ratio, out_channels=self.latent_channels, downscale=True, index_formulas=self.downscale_index_formula, output_device=self.output_device)500 501 def decode(self, samples_in):502 self.throw_exception_if_invalid()503 pixel_samples = None504 try:505 memory_used = self.memory_used_decode(samples_in.shape, self.vae_dtype)506 model_management.load_models_gpu([self.patcher], memory_required=memory_used)507 free_memory = model_management.get_free_memory(self.device)508 batch_number = int(free_memory / memory_used)509 batch_number = max(1, batch_number)510 511 for x in range(0, samples_in.shape[0], batch_number):512 samples = samples_in[x:x+batch_number].to(self.vae_dtype).to(self.device)513 out = self.process_output(self.first_stage_model.decode(samples).to(self.output_device).float())514 if pixel_samples is None:515 pixel_samples = torch.empty((samples_in.shape[0],) + tuple(out.shape[1:]), device=self.output_device)516 pixel_samples[x:x+batch_number] = out517 except model_management.OOM_EXCEPTION:518 logging.warning("Warning: Ran out of memory when regular VAE decoding, retrying with tiled VAE decoding.")519 dims = samples_in.ndim - 2520 if dims == 1:521 pixel_samples = self.decode_tiled_1d(samples_in)522 elif dims == 2:523 pixel_samples = self.decode_tiled_(samples_in)524 elif dims == 3:525 tile = 256 // self.spacial_compression_decode()526 overlap = tile // 4527 pixel_samples = self.decode_tiled_3d(samples_in, tile_x=tile, tile_y=tile, overlap=(1, overlap, overlap))528 529 pixel_samples = pixel_samples.to(self.output_device).movedim(1,-1)530 return pixel_samples531 532 def decode_tiled(self, samples, tile_x=None, tile_y=None, overlap=None, tile_t=None, overlap_t=None):533 self.throw_exception_if_invalid()534 memory_used = self.memory_used_decode(samples.shape, self.vae_dtype) #TODO: calculate mem required for tile535 model_management.load_models_gpu([self.patcher], memory_required=memory_used)536 dims = samples.ndim - 2537 args = {}538 if tile_x is not None:539 args["tile_x"] = tile_x540 if tile_y is not None:541 args["tile_y"] = tile_y542 if overlap is not None:543 args["overlap"] = overlap544 545 if dims == 1:546 args.pop("tile_y")547 output = self.decode_tiled_1d(samples, **args)548 elif dims == 2:549 output = self.decode_tiled_(samples, **args)550 elif dims == 3:551 if overlap_t is None:552 args["overlap"] = (1, overlap, overlap)553 else:554 args["overlap"] = (max(1, overlap_t), overlap, overlap)555 if tile_t is not None:556 args["tile_t"] = max(2, tile_t)557 558 output = self.decode_tiled_3d(samples, **args)559 return output.movedim(1, -1)560 561 def encode(self, pixel_samples):562 self.throw_exception_if_invalid()563 pixel_samples = self.vae_encode_crop_pixels(pixel_samples)564 pixel_samples = pixel_samples.movedim(-1, 1)565 if self.latent_dim == 3 and pixel_samples.ndim < 5:566 pixel_samples = pixel_samples.movedim(1, 0).unsqueeze(0)567 try:568 memory_used = self.memory_used_encode(pixel_samples.shape, self.vae_dtype)569 model_management.load_models_gpu([self.patcher], memory_required=memory_used)570 free_memory = model_management.get_free_memory(self.device)571 batch_number = int(free_memory / max(1, memory_used))572 batch_number = max(1, batch_number)573 samples = None574 for x in range(0, pixel_samples.shape[0], batch_number):575 pixels_in = self.process_input(pixel_samples[x:x + batch_number]).to(self.vae_dtype).to(self.device)576 out = self.first_stage_model.encode(pixels_in).to(self.output_device).float()577 if samples is None:578 samples = torch.empty((pixel_samples.shape[0],) + tuple(out.shape[1:]), device=self.output_device)579 samples[x:x + batch_number] = out580 581 except model_management.OOM_EXCEPTION:582 logging.warning("Warning: Ran out of memory when regular VAE encoding, retrying with tiled VAE encoding.")583 if self.latent_dim == 3:584 tile = 256585 overlap = tile // 4586 samples = self.encode_tiled_3d(pixel_samples, tile_x=tile, tile_y=tile, overlap=(1, overlap, overlap))587 elif self.latent_dim == 1:588 samples = self.encode_tiled_1d(pixel_samples)589 else:590 samples = self.encode_tiled_(pixel_samples)591 592 return samples593 594 def encode_tiled(self, pixel_samples, tile_x=None, tile_y=None, overlap=None, tile_t=None, overlap_t=None):595 self.throw_exception_if_invalid()596 pixel_samples = self.vae_encode_crop_pixels(pixel_samples)597 dims = self.latent_dim598 pixel_samples = pixel_samples.movedim(-1, 1)599 if dims == 3:600 pixel_samples = pixel_samples.movedim(1, 0).unsqueeze(0)601 602 memory_used = self.memory_used_encode(pixel_samples.shape, self.vae_dtype) # TODO: calculate mem required for tile603 model_management.load_models_gpu([self.patcher], memory_required=memory_used)604 605 args = {}606 if tile_x is not None:607 args["tile_x"] = tile_x608 if tile_y is not None:609 args["tile_y"] = tile_y610 if overlap is not None:611 args["overlap"] = overlap612 613 if dims == 1:614 args.pop("tile_y")615 samples = self.encode_tiled_1d(pixel_samples, **args)616 elif dims == 2:617 samples = self.encode_tiled_(pixel_samples, **args)618 elif dims == 3:619 if tile_t is not None:620 tile_t_latent = max(2, self.downscale_ratio[0](tile_t))621 else:622 tile_t_latent = 9999623 args["tile_t"] = self.upscale_ratio[0](tile_t_latent)624 625 if overlap_t is None:626 args["overlap"] = (1, overlap, overlap)627 else:628 args["overlap"] = (self.upscale_ratio[0](max(1, min(tile_t_latent // 2, self.downscale_ratio[0](overlap_t)))), overlap, overlap)629 maximum = pixel_samples.shape[2]630 maximum = self.upscale_ratio[0](self.downscale_ratio[0](maximum))631 632 samples = self.encode_tiled_3d(pixel_samples[:,:,:maximum], **args)633 634 return samples635 636 def get_sd(self):637 return self.first_stage_model.state_dict()638 639 def spacial_compression_decode(self):640 try:641 return self.upscale_ratio[-1]642 except:643 return self.upscale_ratio644 645 def spacial_compression_encode(self):646 try:647 return self.downscale_ratio[-1]648 except:649 return self.downscale_ratio650 651 def temporal_compression_decode(self):652 try:653 return round(self.upscale_ratio[0](8192) / 8192)654 except:655 return None656 657class StyleModel:658 def __init__(self, model, device="cpu"):659 self.model = model660 661 def get_cond(self, input):662 return self.model(input.last_hidden_state)663 664 665def load_style_model(ckpt_path):666 model_data = comfy.utils.load_torch_file(ckpt_path, safe_load=True)667 keys = model_data.keys()668 if "style_embedding" in keys:669 model = comfy.t2i_adapter.adapter.StyleAdapter(width=1024, context_dim=768, num_head=8, n_layes=3, num_token=8)670 elif "redux_down.weight" in keys:671 model = comfy.ldm.flux.redux.ReduxImageEncoder()672 else:673 raise Exception("invalid style model {}".format(ckpt_path))674 model.load_state_dict(model_data)675 return StyleModel(model)676 677class CLIPType(Enum):678 STABLE_DIFFUSION = 1679 STABLE_CASCADE = 2680 SD3 = 3681 STABLE_AUDIO = 4682 HUNYUAN_DIT = 5683 FLUX = 6684 MOCHI = 7685 LTXV = 8686 HUNYUAN_VIDEO = 9687 PIXART = 10688 COSMOS = 11689 LUMINA2 = 12690 WAN = 13691 692 693def load_clip(ckpt_paths, embedding_directory=None, clip_type=CLIPType.STABLE_DIFFUSION, model_options={}):694 clip_data = []695 for p in ckpt_paths:696 clip_data.append(comfy.utils.load_torch_file(p, safe_load=True))697 return load_text_encoder_state_dicts(clip_data, embedding_directory=embedding_directory, clip_type=clip_type, model_options=model_options)698 699 700class TEModel(Enum):701 CLIP_L = 1702 CLIP_H = 2703 CLIP_G = 3704 T5_XXL = 4705 T5_XL = 5706 T5_BASE = 6707 LLAMA3_8 = 7708 T5_XXL_OLD = 8709 GEMMA_2_2B = 9710 711def detect_te_model(sd):712 if "text_model.encoder.layers.30.mlp.fc1.weight" in sd:713 return TEModel.CLIP_G714 if "text_model.encoder.layers.22.mlp.fc1.weight" in sd:715 return TEModel.CLIP_H716 if "text_model.encoder.layers.0.mlp.fc1.weight" in sd:717 return TEModel.CLIP_L718 if "encoder.block.23.layer.1.DenseReluDense.wi_1.weight" in sd:719 weight = sd["encoder.block.23.layer.1.DenseReluDense.wi_1.weight"]720 if weight.shape[-1] == 4096:721 return TEModel.T5_XXL722 elif weight.shape[-1] == 2048:723 return TEModel.T5_XL724 if 'encoder.block.23.layer.1.DenseReluDense.wi.weight' in sd:725 return TEModel.T5_XXL_OLD726 if "encoder.block.0.layer.0.SelfAttention.k.weight" in sd:727 return TEModel.T5_BASE728 if 'model.layers.0.post_feedforward_layernorm.weight' in sd:729 return TEModel.GEMMA_2_2B730 if "model.layers.0.post_attention_layernorm.weight" in sd:731 return TEModel.LLAMA3_8732 return None733 734 735def t5xxl_detect(clip_data):736 weight_name = "encoder.block.23.layer.1.DenseReluDense.wi_1.weight"737 weight_name_old = "encoder.block.23.layer.1.DenseReluDense.wi.weight"738 739 for sd in clip_data:740 if weight_name in sd or weight_name_old in sd:741 return comfy.text_encoders.sd3_clip.t5_xxl_detect(sd)742 743 return {}744 745def llama_detect(clip_data):746 weight_name = "model.layers.0.self_attn.k_proj.weight"747 748 for sd in clip_data:749 if weight_name in sd:750 return comfy.text_encoders.hunyuan_video.llama_detect(sd)751 752 return {}753 754def load_text_encoder_state_dicts(state_dicts=[], embedding_directory=None, clip_type=CLIPType.STABLE_DIFFUSION, model_options={}):755 clip_data = state_dicts756 757 class EmptyClass:758 pass759 760 for i in range(len(clip_data)):761 if "transformer.resblocks.0.ln_1.weight" in clip_data[i]:762 clip_data[i] = comfy.utils.clip_text_transformers_convert(clip_data[i], "", "")763 else:764 if "text_projection" in clip_data[i]:765 clip_data[i]["text_projection.weight"] = clip_data[i]["text_projection"].transpose(0, 1) #old models saved with the CLIPSave node766 767 tokenizer_data = {}768 clip_target = EmptyClass()769 clip_target.params = {}770 if len(clip_data) == 1:771 te_model = detect_te_model(clip_data[0])772 if te_model == TEModel.CLIP_G:773 if clip_type == CLIPType.STABLE_CASCADE:774 clip_target.clip = sdxl_clip.StableCascadeClipModel775 clip_target.tokenizer = sdxl_clip.StableCascadeTokenizer776 elif clip_type == CLIPType.SD3:777 clip_target.clip = comfy.text_encoders.sd3_clip.sd3_clip(clip_l=False, clip_g=True, t5=False)778 clip_target.tokenizer = comfy.text_encoders.sd3_clip.SD3Tokenizer779 else:780 clip_target.clip = sdxl_clip.SDXLRefinerClipModel781 clip_target.tokenizer = sdxl_clip.SDXLTokenizer782 elif te_model == TEModel.CLIP_H:783 clip_target.clip = comfy.text_encoders.sd2_clip.SD2ClipModel784 clip_target.tokenizer = comfy.text_encoders.sd2_clip.SD2Tokenizer785 elif te_model == TEModel.T5_XXL:786 if clip_type == CLIPType.SD3:787 clip_target.clip = comfy.text_encoders.sd3_clip.sd3_clip(clip_l=False, clip_g=False, t5=True, **t5xxl_detect(clip_data))788 clip_target.tokenizer = comfy.text_encoders.sd3_clip.SD3Tokenizer789 elif clip_type == CLIPType.LTXV:790 clip_target.clip = comfy.text_encoders.lt.ltxv_te(**t5xxl_detect(clip_data))791 clip_target.tokenizer = comfy.text_encoders.lt.LTXVT5Tokenizer792 elif clip_type == CLIPType.PIXART:793 clip_target.clip = comfy.text_encoders.pixart_t5.pixart_te(**t5xxl_detect(clip_data))794 clip_target.tokenizer = comfy.text_encoders.pixart_t5.PixArtTokenizer795 elif clip_type == CLIPType.WAN:796 clip_target.clip = comfy.text_encoders.wan.te(**t5xxl_detect(clip_data))797 clip_target.tokenizer = comfy.text_encoders.wan.WanT5Tokenizer798 tokenizer_data["spiece_model"] = clip_data[0].get("spiece_model", None)799 else: #CLIPType.MOCHI800 clip_target.clip = comfy.text_encoders.genmo.mochi_te(**t5xxl_detect(clip_data))801 clip_target.tokenizer = comfy.text_encoders.genmo.MochiT5Tokenizer802 elif te_model == TEModel.T5_XXL_OLD:803 clip_target.clip = comfy.text_encoders.cosmos.te(**t5xxl_detect(clip_data))804 clip_target.tokenizer = comfy.text_encoders.cosmos.CosmosT5Tokenizer805 elif te_model == TEModel.T5_XL:806 clip_target.clip = comfy.text_encoders.aura_t5.AuraT5Model807 clip_target.tokenizer = comfy.text_encoders.aura_t5.AuraT5Tokenizer808 elif te_model == TEModel.T5_BASE:809 clip_target.clip = comfy.text_encoders.sa_t5.SAT5Model810 clip_target.tokenizer = comfy.text_encoders.sa_t5.SAT5Tokenizer811 elif te_model == TEModel.GEMMA_2_2B:812 clip_target.clip = comfy.text_encoders.lumina2.te(**llama_detect(clip_data))813 clip_target.tokenizer = comfy.text_encoders.lumina2.LuminaTokenizer814 tokenizer_data["spiece_model"] = clip_data[0].get("spiece_model", None)815 else:816 if clip_type == CLIPType.SD3:817 clip_target.clip = comfy.text_encoders.sd3_clip.sd3_clip(clip_l=True, clip_g=False, t5=False)818 clip_target.tokenizer = comfy.text_encoders.sd3_clip.SD3Tokenizer819 else:820 clip_target.clip = sd1_clip.SD1ClipModel821 clip_target.tokenizer = sd1_clip.SD1Tokenizer822 elif len(clip_data) == 2:823 if clip_type == CLIPType.SD3:824 te_models = [detect_te_model(clip_data[0]), detect_te_model(clip_data[1])]825 clip_target.clip = comfy.text_encoders.sd3_clip.sd3_clip(clip_l=TEModel.CLIP_L in te_models, clip_g=TEModel.CLIP_G in te_models, t5=TEModel.T5_XXL in te_models, **t5xxl_detect(clip_data))826 clip_target.tokenizer = comfy.text_encoders.sd3_clip.SD3Tokenizer827 elif clip_type == CLIPType.HUNYUAN_DIT:828 clip_target.clip = comfy.text_encoders.hydit.HyditModel829 clip_target.tokenizer = comfy.text_encoders.hydit.HyditTokenizer830 elif clip_type == CLIPType.FLUX:831 clip_target.clip = comfy.text_encoders.flux.flux_clip(**t5xxl_detect(clip_data))832 clip_target.tokenizer = comfy.text_encoders.flux.FluxTokenizer833 elif clip_type == CLIPType.HUNYUAN_VIDEO:834 clip_target.clip = comfy.text_encoders.hunyuan_video.hunyuan_video_clip(**llama_detect(clip_data))835 clip_target.tokenizer = comfy.text_encoders.hunyuan_video.HunyuanVideoTokenizer836 else:837 clip_target.clip = sdxl_clip.SDXLClipModel838 clip_target.tokenizer = sdxl_clip.SDXLTokenizer839 elif len(clip_data) == 3:840 clip_target.clip = comfy.text_encoders.sd3_clip.sd3_clip(**t5xxl_detect(clip_data))841 clip_target.tokenizer = comfy.text_encoders.sd3_clip.SD3Tokenizer842 843 parameters = 0844 for c in clip_data:845 parameters += comfy.utils.calculate_parameters(c)846 tokenizer_data, model_options = comfy.text_encoders.long_clipl.model_options_long_clip(c, tokenizer_data, model_options)847 848 clip = CLIP(clip_target, embedding_directory=embedding_directory, parameters=parameters, tokenizer_data=tokenizer_data, model_options=model_options)849 for c in clip_data:850 m, u = clip.load_sd(c)851 if len(m) > 0:852 logging.warning("clip missing: {}".format(m))853 854 if len(u) > 0:855 logging.debug("clip unexpected: {}".format(u))856 return clip857 858def load_gligen(ckpt_path):859 data = comfy.utils.load_torch_file(ckpt_path, safe_load=True)860 model = gligen.load_gligen(data)861 if model_management.should_use_fp16():862 model = model.half()863 return comfy.model_patcher.ModelPatcher(model, load_device=model_management.get_torch_device(), offload_device=model_management.unet_offload_device())864 865def load_checkpoint(config_path=None, ckpt_path=None, output_vae=True, output_clip=True, embedding_directory=None, state_dict=None, config=None):866 logging.warning("Warning: The load checkpoint with config function is deprecated and will eventually be removed, please use the other one.")867 model, clip, vae, _ = load_checkpoint_guess_config(ckpt_path, output_vae=output_vae, output_clip=output_clip, output_clipvision=False, embedding_directory=embedding_directory, output_model=True)868 #TODO: this function is a mess and should be removed eventually869 if config is None:870 with open(config_path, 'r') as stream:871 config = yaml.safe_load(stream)872 model_config_params = config['model']['params']873 clip_config = model_config_params['cond_stage_config']874 875 if "parameterization" in model_config_params:876 if model_config_params["parameterization"] == "v":877 m = model.clone()878 class ModelSamplingAdvanced(comfy.model_sampling.ModelSamplingDiscrete, comfy.model_sampling.V_PREDICTION):879 pass880 m.add_object_patch("model_sampling", ModelSamplingAdvanced(model.model.model_config))881 model = m882 883 layer_idx = clip_config.get("params", {}).get("layer_idx", None)884 if layer_idx is not None:885 clip.clip_layer(layer_idx)886 887 return (model, clip, vae)888 889def load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, output_clipvision=False, embedding_directory=None, output_model=True, model_options={}, te_model_options={}):890 sd, metadata = comfy.utils.load_torch_file(ckpt_path, return_metadata=True)891 out = load_state_dict_guess_config(sd, output_vae, output_clip, output_clipvision, embedding_directory, output_model, model_options, te_model_options=te_model_options, metadata=metadata)892 if out is None:893 raise RuntimeError("ERROR: Could not detect model type of: {}".format(ckpt_path))894 return out895 896def load_state_dict_guess_config(sd, output_vae=True, output_clip=True, output_clipvision=False, embedding_directory=None, output_model=True, model_options={}, te_model_options={}, metadata=None):897 clip = None898 clipvision = None899 vae = None900 model = None901 model_patcher = None902 903 diffusion_model_prefix = model_detection.unet_prefix_from_state_dict(sd)904 parameters = comfy.utils.calculate_parameters(sd, diffusion_model_prefix)905 weight_dtype = comfy.utils.weight_dtype(sd, diffusion_model_prefix)906 load_device = model_management.get_torch_device()907 908 model_config = model_detection.model_config_from_unet(sd, diffusion_model_prefix, metadata=metadata)909 if model_config is None:910 logging.warning("Warning, This is not a checkpoint file, trying to load it as a diffusion model only.")911 diffusion_model = load_diffusion_model_state_dict(sd, model_options={})912 if diffusion_model is None:913 return None914 return (diffusion_model, None, VAE(sd={}), None) # The VAE object is there to throw an exception if it's actually used'915 916 917 unet_weight_dtype = list(model_config.supported_inference_dtypes)918 if model_config.scaled_fp8 is not None:919 weight_dtype = None920 921 model_config.custom_operations = model_options.get("custom_operations", None)922 unet_dtype = model_options.get("dtype", model_options.get("weight_dtype", None))923 924 if unet_dtype is None:925 unet_dtype = model_management.unet_dtype(model_params=parameters, supported_dtypes=unet_weight_dtype, weight_dtype=weight_dtype)926 927 manual_cast_dtype = model_management.unet_manual_cast(unet_dtype, load_device, model_config.supported_inference_dtypes)928 model_config.set_inference_dtype(unet_dtype, manual_cast_dtype)929 930 if model_config.clip_vision_prefix is not None:931 if output_clipvision:932 clipvision = clip_vision.load_clipvision_from_sd(sd, model_config.clip_vision_prefix, True)933 934 if output_model:935 inital_load_device = model_management.unet_inital_load_device(parameters, unet_dtype)936 model = model_config.get_model(sd, diffusion_model_prefix, device=inital_load_device)937 model.load_model_weights(sd, diffusion_model_prefix)938 939 if output_vae:940 vae_sd = comfy.utils.state_dict_prefix_replace(sd, {k: "" for k in model_config.vae_key_prefix}, filter_keys=True)941 vae_sd = model_config.process_vae_state_dict(vae_sd)942 vae = VAE(sd=vae_sd, metadata=metadata)943 944 if output_clip:945 clip_target = model_config.clip_target(state_dict=sd)946 if clip_target is not None:947 clip_sd = model_config.process_clip_state_dict(sd)948 if len(clip_sd) > 0:949 parameters = comfy.utils.calculate_parameters(clip_sd)950 clip = CLIP(clip_target, embedding_directory=embedding_directory, tokenizer_data=clip_sd, parameters=parameters, model_options=te_model_options)951 m, u = clip.load_sd(clip_sd, full_model=True)952 if len(m) > 0:953 m_filter = list(filter(lambda a: ".logit_scale" not in a and ".transformer.text_projection.weight" not in a, m))954 if len(m_filter) > 0:955 logging.warning("clip missing: {}".format(m))956 else:957 logging.debug("clip missing: {}".format(m))958 959 if len(u) > 0:960 logging.debug("clip unexpected {}:".format(u))961 else:962 logging.warning("no CLIP/text encoder weights in checkpoint, the text encoder model will not be loaded.")963 964 left_over = sd.keys()965 if len(left_over) > 0:966 logging.debug("left over keys: {}".format(left_over))967 968 if output_model:969 model_patcher = comfy.model_patcher.ModelPatcher(model, load_device=load_device, offload_device=model_management.unet_offload_device())970 if inital_load_device != torch.device("cpu"):971 logging.info("loaded diffusion model directly to GPU")972 model_management.load_models_gpu([model_patcher], force_full_load=True)973 974 return (model_patcher, clip, vae, clipvision)975 976 977def load_diffusion_model_state_dict(sd, model_options={}): #load unet in diffusers or regular format978 dtype = model_options.get("dtype", None)979 980 #Allow loading unets from checkpoint files981 diffusion_model_prefix = model_detection.unet_prefix_from_state_dict(sd)982 temp_sd = comfy.utils.state_dict_prefix_replace(sd, {diffusion_model_prefix: ""}, filter_keys=True)983 if len(temp_sd) > 0:984 sd = temp_sd985 986 parameters = comfy.utils.calculate_parameters(sd)987 weight_dtype = comfy.utils.weight_dtype(sd)988 989 load_device = model_management.get_torch_device()990 model_config = model_detection.model_config_from_unet(sd, "")991 992 if model_config is not None:993 new_sd = sd994 else:995 new_sd = model_detection.convert_diffusers_mmdit(sd, "")996 if new_sd is not None: #diffusers mmdit997 model_config = model_detection.model_config_from_unet(new_sd, "")998 if model_config is None:999 return None1000 else: #diffusers unet1001 model_config = model_detection.model_config_from_diffusers_unet(sd)1002 if model_config is None:1003 return None1004 1005 diffusers_keys = comfy.utils.unet_to_diffusers(model_config.unet_config)1006 1007 new_sd = {}1008 for k in diffusers_keys:1009 if k in sd:1010 new_sd[diffusers_keys[k]] = sd.pop(k)1011 else:1012 logging.warning("{} {}".format(diffusers_keys[k], k))1013 1014 offload_device = model_management.unet_offload_device()1015 unet_weight_dtype = list(model_config.supported_inference_dtypes)1016 if model_config.scaled_fp8 is not None:1017 weight_dtype = None1018 1019 if dtype is None:1020 unet_dtype = model_management.unet_dtype(model_params=parameters, supported_dtypes=unet_weight_dtype, weight_dtype=weight_dtype)1021 else:1022 unet_dtype = dtype1023 1024 manual_cast_dtype = model_management.unet_manual_cast(unet_dtype, load_device, model_config.supported_inference_dtypes)1025 model_config.set_inference_dtype(unet_dtype, manual_cast_dtype)1026 model_config.custom_operations = model_options.get("custom_operations", model_config.custom_operations)1027 if model_options.get("fp8_optimizations", False):1028 model_config.optimizations["fp8"] = True1029 1030 model = model_config.get_model(new_sd, "")1031 model = model.to(offload_device)1032 model.load_model_weights(new_sd, "")1033 left_over = sd.keys()1034 if len(left_over) > 0:1035 logging.info("left over keys in unet: {}".format(left_over))1036 return comfy.model_patcher.ModelPatcher(model, load_device=load_device, offload_device=offload_device)1037 1038 1039def load_diffusion_model(unet_path, model_options={}):1040 sd = comfy.utils.load_torch_file(unet_path)1041 model = load_diffusion_model_state_dict(sd, model_options=model_options)1042 if model is None:1043 logging.error("ERROR UNSUPPORTED UNET {}".format(unet_path))1044 raise RuntimeError("ERROR: Could not detect model type of: {}".format(unet_path))1045 return model1046 1047def load_unet(unet_path, dtype=None):1048 logging.warning("The load_unet function has been deprecated and will be removed please switch to: load_diffusion_model")1049 return load_diffusion_model(unet_path, model_options={"dtype": dtype})1050 1051def load_unet_state_dict(sd, dtype=None):1052 logging.warning("The load_unet_state_dict function has been deprecated and will be removed please switch to: load_diffusion_model_state_dict")1053 return load_diffusion_model_state_dict(sd, model_options={"dtype": dtype})1054 1055def save_checkpoint(output_path, model, clip=None, vae=None, clip_vision=None, metadata=None, extra_keys={}):1056 clip_sd = None1057 load_models = [model]1058 if clip is not None:1059 load_models.append(clip.load_model())1060 clip_sd = clip.get_sd()1061 vae_sd = None1062 if vae is not None:1063 vae_sd = vae.get_sd()1064 1065 model_management.load_models_gpu(load_models, force_patch_weights=True)1066 clip_vision_sd = clip_vision.get_sd() if clip_vision is not None else None1067 sd = model.model.state_dict_for_saving(clip_sd, vae_sd, clip_vision_sd)1068 for k in extra_keys:1069 sd[k] = extra_keys[k]1070 1071 for k in sd:1072 t = sd[k]1073 if not t.is_contiguous():1074 sd[k] = t.contiguous()1075 1076 comfy.utils.save_torch_file(sd, output_path, metadata=metadata)1077