Support fp16 for Cosmos-Predict2 and Anima (#12249)

This commit is contained in:
tdrussell
2026-02-06 19:12:15 -06:00
committed by GitHub
parent 204e65b8dc
commit 6a26328842
2 changed files with 19 additions and 9 deletions
+2 -2
View File
@@ -993,7 +993,7 @@ class CosmosT2IPredict2(supported_models_base.BASE):
memory_usage_factor = 1.0
supported_inference_dtypes = [torch.bfloat16, torch.float32]
supported_inference_dtypes = [torch.bfloat16, torch.float16, torch.float32]
def __init__(self, unet_config):
super().__init__(unet_config)
@@ -1023,7 +1023,7 @@ class Anima(supported_models_base.BASE):
memory_usage_factor = 1.0
supported_inference_dtypes = [torch.bfloat16, torch.float32]
supported_inference_dtypes = [torch.bfloat16, torch.float16, torch.float32]
def __init__(self, unet_config):
super().__init__(unet_config)