Mightys/Notebook_Scripts

1# A very basic ComfyUI plugin to patch anima, which is based cosmos pt2,2# to use fp16 "safely" on old GPUs.3# Only tested for anima.4# Written by reakaakasky (https://civitai.com/user/reakaakasky).5 6# To clarify, this plugin isn't uploaded to GitHub because it's a very "dirty" plugin.7# "dirty" means it hot-patches/monkey patches the comfyui core code.8# This approach is terrible from a programming perspective.9# But this is the simplest approach I can think of.10 11# How to use:12# Put this file in the ComfyUI "custom_nodes" dir.13# Use "ModelComputeDtype" node and set dtype to "fp16".14# To disable the patch, remove the file, or rename the ".py" suffix, to something15# like ".disable", whatever.16 17# Version: v1 (2/2/2026)18# Version: v1.1 (2/2/2026) optimazed for torch.compile19 20import torch21import logging22 23 24logger = logging.getLogger(__name__)25logger.info("[anima fp16 patch] patch loading")26 27NODE_CLASS_MAPPINGS = {}28 29NODE_DISPLAY_NAME_MAPPINGS = {}30 31import comfy.ldm.cosmos.predict2 as p232 33ampf16 = torch.autocast("cuda", dtype=torch.float16)34ampf32 = torch.autocast("cuda", dtype=torch.float32)35 36 37def p2_Block_init_patch(self: p2.Block, *args, **kwargs):38    self.__init_org(*args, **kwargs)39 40    self.adaln_modulation_self_attn.forward = ampf16(self.adaln_modulation_self_attn.forward)41    self.adaln_modulation_cross_attn.forward = ampf16(self.adaln_modulation_cross_attn.forward)42 43    self.self_attn.forward = ampf16(self.self_attn.forward)44    self.cross_attn.forward = ampf16(self.cross_attn.forward)45    self.mlp.forward = ampf16(self.mlp.forward)46 47    self.forward = ampf32(self.forward)48 49 50p2.Block.__init_org = p2.Block.__init__51p2.Block.__init__ = p2_Block_init_patch52 53torch.set_float32_matmul_precision("high")54torch.backends.cuda.matmul.allow_fp16_accumulation = True55 56logger.info("[anima fp16 patch] patch loaded")57