Merge pull request #28 from Lightricks/feature/model_dtype

Add model dtype selector.

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(),
         )

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:

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}

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(),
         )