diff --git a/loader_node.py b/loader_node.py index 074dc85..496e7bb 100644 --- a/loader_node.py +++ b/loader_node.py @@ -29,6 +29,7 @@ class LTXVLoader: folder_paths.get_filename_list("checkpoints"), {"tooltip": "The name of the checkpoint (model) to load."}, ), + "dtype": (["bfloat16", "float32"], {"default": "float32"}), } } @@ -39,7 +40,8 @@ class LTXVLoader: TITLE = "LTXV Loader" OUTPUT_NODE = False - def load(self, ckpt_name): + def load(self, ckpt_name, dtype): + dtype_map = {"bfloat16": torch.bfloat16, "float32": torch.float32} load_device = comfy.model_management.get_torch_device() offload_device = comfy.model_management.unet_offload_device() @@ -51,7 +53,7 @@ class LTXVLoader: num_latent_channels = vae.first_stage_model.config.latent_channels model = self._load_unet( - load_device, offload_device, weights, num_latent_channels + load_device, offload_device, weights, num_latent_channels, dtype=dtype_map[dtype] ) return (model, vae) @@ -92,7 +94,7 @@ class LTXVLoader: ) return vae - def _load_unet(self, load_device, offload_device, weights, num_latent_channels): + def _load_unet(self, load_device, offload_device, weights, num_latent_channels, dtype): config = { "_class_name": "Transformer3DModel", "_diffusers_version": "0.25.1", @@ -135,11 +137,11 @@ class LTXVLoader: if key.startswith(unet_prefix) } ) - transformer.to(torch.float32).to(load_device).eval() + transformer.to(dtype).to(load_device).eval() diffusion_model = LTXVTransformer3D(transformer, SymmetricPatchifier(1)) model = LTXVModel( - LTXVModelConfig(num_latent_channels), + LTXVModelConfig(num_latent_channels, dtype=dtype), model_type=comfy.model_base.ModelType.FLOW, device=comfy.model_management.get_torch_device(), ) diff --git a/model.py b/model.py index cb8c288..555d26b 100644 --- a/model.py +++ b/model.py @@ -18,14 +18,14 @@ from ltx_video.models.autoencoders.vae_encode import get_vae_size_scale_factor class LTXVModelConfig: - def __init__(self, latent_channels): + def __init__(self, latent_channels, dtype): self.unet_config = {} self.unet_extra_config = {} self.latent_format = comfy.latent_formats.LatentFormat() self.latent_format.latent_channels = latent_channels - self.manual_cast_dtype = torch.float32 + self.manual_cast_dtype = dtype self.sampling_settings = {"multiplier": 1.0} - self.memory_usage_factor = 1 + self.memory_usage_factor = 2.7 # denoiser is handled by extension self.unet_config["disable_unet_model_creation"] = True @@ -107,7 +107,7 @@ class LTXVTransformer3D(nn.Module): # infer mask from context padding, assumes padding vectors are all zero. latent = latent.to(self.transformer.dtype) latent_patchified = self.patchifier.patchify(latent) - context_mask = (context != 0).any(dim=2).float() + context_mask = (context != 0).any(dim=2).to(self.transformer.dtype) l = latent_patchified if mixed_precision: diff --git a/t5_encoder.py b/t5_encoder.py index 3ffcb05..9ba6542 100644 --- a/t5_encoder.py +++ b/t5_encoder.py @@ -84,6 +84,7 @@ class LTXVTextEncoderModel(torch.nn.Module): [[w[1] for w in token_weight_pairs_t5]], device=self.t5xxl.device, ) + self.to(self.t5xxl.device) # comfyui skips loading some weights to gpu out = self.t5xxl(text_input_ids, attention_mask=prompt_attention_mask)[0] out = out * prompt_attention_mask.unsqueeze(2) return out, None, {"attention_mask": prompt_attention_mask} diff --git a/transformer.py b/transformer.py index 27a1bff..f389a3d 100644 --- a/transformer.py +++ b/transformer.py @@ -206,7 +206,7 @@ class LTXVModelConfigurator: ) model.model.diffusion_model.to(load_device) new_model = LTXVModel( - LTXVModelConfig(vae.first_stage_model.config.latent_channels), + LTXVModelConfig(vae.first_stage_model.config.latent_channels, model.model.model_config.manual_cast_dtype), model_type=comfy.model_base.ModelType.FLOW, device=comfy.model_management.get_torch_device(), )