mirror of
https://git.datalinker.icu/comfyanonymous/ComfyUI
synced 2026-11-12 12:40:04 +08:00
87 lines
2.8 KiB
Python
87 lines
2.8 KiB
Python
import torch
|
|
from torch.autograd import Variable
|
|
from torch.nn.parameter import Parameter
|
|
|
|
|
|
_FLOAT_TYPES = (torch.FloatTensor, torch.cuda.FloatTensor)
|
|
_HALF_TYPES = (torch.HalfTensor, torch.cuda.HalfTensor)
|
|
|
|
|
|
def conversion_helper(val, conversion):
|
|
"""Apply conversion to val. Recursively apply conversion if `val`
|
|
#is a nested tuple/list structure."""
|
|
if isinstance(val, dict):
|
|
res_dict = {}
|
|
for k, v in val.items():
|
|
if k!= 'cos_cis_img' and k != 'sin_cis_img':
|
|
res_dict[k] = conversion_helper(v, conversion)
|
|
else:
|
|
res_dict[k] = v
|
|
return res_dict
|
|
if not isinstance(val, (tuple, list)):
|
|
return conversion(val)
|
|
rtn = [conversion_helper(v, conversion) for v in val]
|
|
if isinstance(val, tuple):
|
|
rtn = tuple(rtn)
|
|
return rtn
|
|
|
|
|
|
def fp32_to_float16(val, float16_convertor):
|
|
"""Convert fp32 `val` to fp16/bf16"""
|
|
def half_conversion(val):
|
|
val_typecheck = val
|
|
if isinstance(val_typecheck, (Parameter, Variable)):
|
|
val_typecheck = val.data
|
|
if isinstance(val_typecheck, _FLOAT_TYPES):
|
|
val = float16_convertor(val)
|
|
return val
|
|
return conversion_helper(val, half_conversion)
|
|
|
|
|
|
def float16_to_fp32(val):
|
|
"""Convert fp16/bf16 `val` to fp32"""
|
|
def float_conversion(val):
|
|
val_typecheck = val
|
|
if isinstance(val_typecheck, (Parameter, Variable)):
|
|
val_typecheck = val.data
|
|
if isinstance(val_typecheck, (_HALF_TYPES,)):
|
|
val = val.float()
|
|
return val
|
|
return conversion_helper(val, float_conversion)
|
|
|
|
|
|
class Float16Module(torch.nn.Module):
|
|
|
|
def __init__(self, module, args):
|
|
super(Float16Module, self).__init__()
|
|
|
|
self.add_module('module', module.half())
|
|
|
|
def float16_convertor(val):
|
|
return val.half()
|
|
|
|
self.float16_convertor = float16_convertor
|
|
|
|
self.config = self.module.config
|
|
self.dtype = torch.float16
|
|
|
|
def forward(self, *inputs, **kwargs):
|
|
inputs = fp32_to_float16(inputs, self.float16_convertor)
|
|
kwargs = fp32_to_float16(kwargs, self.float16_convertor)
|
|
outputs = self.module(*inputs, **kwargs)
|
|
outputs = float16_to_fp32(outputs)
|
|
return outputs
|
|
|
|
def state_dict(self, destination=None, prefix='', keep_vars=False):
|
|
return self.module.state_dict(destination, prefix, keep_vars)
|
|
|
|
|
|
def state_dict_for_save_checkpoint(self, destination=None, prefix='',
|
|
keep_vars=False):
|
|
return self.module.state_dict_for_save_checkpoint(destination, prefix,
|
|
keep_vars)
|
|
|
|
def load_state_dict(self, state_dict, strict=True):
|
|
self.module.load_state_dict(state_dict, strict=strict)
|
|
|