fred-dev/comfy_ui_ali

1"""2    This file is part of ComfyUI.3    Copyright (C) 2024 Stability AI4 5    This program is free software: you can redistribute it and/or modify6    it under the terms of the GNU General Public License as published by7    the Free Software Foundation, either version 3 of the License, or8    (at your option) any later version.9 10    This program is distributed in the hope that it will be useful,11    but WITHOUT ANY WARRANTY; without even the implied warranty of12    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the13    GNU General Public License for more details.14 15    You should have received a copy of the GNU General Public License16    along with this program.  If not, see <https://www.gnu.org/licenses/>.17"""18 19import torch20import logging21import comfy.model_management22from comfy.cli_args import args, PerformanceFeature23import comfy.float24 25cast_to = comfy.model_management.cast_to #TODO: remove once no more references26 27def cast_to_input(weight, input, non_blocking=False, copy=True):28    return comfy.model_management.cast_to(weight, input.dtype, input.device, non_blocking=non_blocking, copy=copy)29 30def cast_bias_weight(s, input=None, dtype=None, device=None, bias_dtype=None):31    if input is not None:32        if dtype is None:33            dtype = input.dtype34        if bias_dtype is None:35            bias_dtype = dtype36        if device is None:37            device = input.device38 39    bias = None40    non_blocking = comfy.model_management.device_supports_non_blocking(device)41    if s.bias is not None:42        has_function = len(s.bias_function) > 043        bias = comfy.model_management.cast_to(s.bias, bias_dtype, device, non_blocking=non_blocking, copy=has_function)44        if has_function:45            for f in s.bias_function:46                bias = f(bias)47 48    has_function = len(s.weight_function) > 049    weight = comfy.model_management.cast_to(s.weight, dtype, device, non_blocking=non_blocking, copy=has_function)50    if has_function:51        for f in s.weight_function:52            weight = f(weight)53    return weight, bias54 55class CastWeightBiasOp:56    comfy_cast_weights = False57    weight_function = []58    bias_function = []59 60class disable_weight_init:61    class Linear(torch.nn.Linear, CastWeightBiasOp):62        def reset_parameters(self):63            return None64 65        def forward_comfy_cast_weights(self, input):66            weight, bias = cast_bias_weight(self, input)67            return torch.nn.functional.linear(input, weight, bias)68 69        def forward(self, *args, **kwargs):70            if self.comfy_cast_weights or len(self.weight_function) > 0 or len(self.bias_function) > 0:71                return self.forward_comfy_cast_weights(*args, **kwargs)72            else:73                return super().forward(*args, **kwargs)74 75    class Conv1d(torch.nn.Conv1d, CastWeightBiasOp):76        def reset_parameters(self):77            return None78 79        def forward_comfy_cast_weights(self, input):80            weight, bias = cast_bias_weight(self, input)81            return self._conv_forward(input, weight, bias)82 83        def forward(self, *args, **kwargs):84            if self.comfy_cast_weights or len(self.weight_function) > 0 or len(self.bias_function) > 0:85                return self.forward_comfy_cast_weights(*args, **kwargs)86            else:87                return super().forward(*args, **kwargs)88 89    class Conv2d(torch.nn.Conv2d, CastWeightBiasOp):90        def reset_parameters(self):91            return None92 93        def forward_comfy_cast_weights(self, input):94            weight, bias = cast_bias_weight(self, input)95            return self._conv_forward(input, weight, bias)96 97        def forward(self, *args, **kwargs):98            if self.comfy_cast_weights or len(self.weight_function) > 0 or len(self.bias_function) > 0:99                return self.forward_comfy_cast_weights(*args, **kwargs)100            else:101                return super().forward(*args, **kwargs)102 103    class Conv3d(torch.nn.Conv3d, CastWeightBiasOp):104        def reset_parameters(self):105            return None106 107        def forward_comfy_cast_weights(self, input):108            weight, bias = cast_bias_weight(self, input)109            return self._conv_forward(input, weight, bias)110 111        def forward(self, *args, **kwargs):112            if self.comfy_cast_weights or len(self.weight_function) > 0 or len(self.bias_function) > 0:113                return self.forward_comfy_cast_weights(*args, **kwargs)114            else:115                return super().forward(*args, **kwargs)116 117    class GroupNorm(torch.nn.GroupNorm, CastWeightBiasOp):118        def reset_parameters(self):119            return None120 121        def forward_comfy_cast_weights(self, input):122            weight, bias = cast_bias_weight(self, input)123            return torch.nn.functional.group_norm(input, self.num_groups, weight, bias, self.eps)124 125        def forward(self, *args, **kwargs):126            if self.comfy_cast_weights or len(self.weight_function) > 0 or len(self.bias_function) > 0:127                return self.forward_comfy_cast_weights(*args, **kwargs)128            else:129                return super().forward(*args, **kwargs)130 131    class LayerNorm(torch.nn.LayerNorm, CastWeightBiasOp):132        def reset_parameters(self):133            return None134 135        def forward_comfy_cast_weights(self, input):136            if self.weight is not None:137                weight, bias = cast_bias_weight(self, input)138            else:139                weight = None140                bias = None141            return torch.nn.functional.layer_norm(input, self.normalized_shape, weight, bias, self.eps)142 143        def forward(self, *args, **kwargs):144            if self.comfy_cast_weights or len(self.weight_function) > 0 or len(self.bias_function) > 0:145                return self.forward_comfy_cast_weights(*args, **kwargs)146            else:147                return super().forward(*args, **kwargs)148 149    class ConvTranspose2d(torch.nn.ConvTranspose2d, CastWeightBiasOp):150        def reset_parameters(self):151            return None152 153        def forward_comfy_cast_weights(self, input, output_size=None):154            num_spatial_dims = 2155            output_padding = self._output_padding(156                input, output_size, self.stride, self.padding, self.kernel_size,157                num_spatial_dims, self.dilation)158 159            weight, bias = cast_bias_weight(self, input)160            return torch.nn.functional.conv_transpose2d(161                input, weight, bias, self.stride, self.padding,162                output_padding, self.groups, self.dilation)163 164        def forward(self, *args, **kwargs):165            if self.comfy_cast_weights or len(self.weight_function) > 0 or len(self.bias_function) > 0:166                return self.forward_comfy_cast_weights(*args, **kwargs)167            else:168                return super().forward(*args, **kwargs)169 170    class ConvTranspose1d(torch.nn.ConvTranspose1d, CastWeightBiasOp):171        def reset_parameters(self):172            return None173 174        def forward_comfy_cast_weights(self, input, output_size=None):175            num_spatial_dims = 1176            output_padding = self._output_padding(177                input, output_size, self.stride, self.padding, self.kernel_size,178                num_spatial_dims, self.dilation)179 180            weight, bias = cast_bias_weight(self, input)181            return torch.nn.functional.conv_transpose1d(182                input, weight, bias, self.stride, self.padding,183                output_padding, self.groups, self.dilation)184 185        def forward(self, *args, **kwargs):186            if self.comfy_cast_weights or len(self.weight_function) > 0 or len(self.bias_function) > 0:187                return self.forward_comfy_cast_weights(*args, **kwargs)188            else:189                return super().forward(*args, **kwargs)190 191    class Embedding(torch.nn.Embedding, CastWeightBiasOp):192        def reset_parameters(self):193            self.bias = None194            return None195 196        def forward_comfy_cast_weights(self, input, out_dtype=None):197            output_dtype = out_dtype198            if self.weight.dtype == torch.float16 or self.weight.dtype == torch.bfloat16:199                out_dtype = None200            weight, bias = cast_bias_weight(self, device=input.device, dtype=out_dtype)201            return torch.nn.functional.embedding(input, weight, self.padding_idx, self.max_norm, self.norm_type, self.scale_grad_by_freq, self.sparse).to(dtype=output_dtype)202 203        def forward(self, *args, **kwargs):204            if self.comfy_cast_weights or len(self.weight_function) > 0 or len(self.bias_function) > 0:205                return self.forward_comfy_cast_weights(*args, **kwargs)206            else:207                if "out_dtype" in kwargs:208                    kwargs.pop("out_dtype")209                return super().forward(*args, **kwargs)210 211    @classmethod212    def conv_nd(s, dims, *args, **kwargs):213        if dims == 2:214            return s.Conv2d(*args, **kwargs)215        elif dims == 3:216            return s.Conv3d(*args, **kwargs)217        else:218            raise ValueError(f"unsupported dimensions: {dims}")219 220 221class manual_cast(disable_weight_init):222    class Linear(disable_weight_init.Linear):223        comfy_cast_weights = True224 225    class Conv1d(disable_weight_init.Conv1d):226        comfy_cast_weights = True227 228    class Conv2d(disable_weight_init.Conv2d):229        comfy_cast_weights = True230 231    class Conv3d(disable_weight_init.Conv3d):232        comfy_cast_weights = True233 234    class GroupNorm(disable_weight_init.GroupNorm):235        comfy_cast_weights = True236 237    class LayerNorm(disable_weight_init.LayerNorm):238        comfy_cast_weights = True239 240    class ConvTranspose2d(disable_weight_init.ConvTranspose2d):241        comfy_cast_weights = True242 243    class ConvTranspose1d(disable_weight_init.ConvTranspose1d):244        comfy_cast_weights = True245 246    class Embedding(disable_weight_init.Embedding):247        comfy_cast_weights = True248 249 250def fp8_linear(self, input):251    dtype = self.weight.dtype252    if dtype not in [torch.float8_e4m3fn]:253        return None254 255    tensor_2d = False256    if len(input.shape) == 2:257        tensor_2d = True258        input = input.unsqueeze(1)259 260    input_shape = input.shape261    input_dtype = input.dtype262    if len(input.shape) == 3:263        w, bias = cast_bias_weight(self, input, dtype=dtype, bias_dtype=input_dtype)264        w = w.t()265 266        scale_weight = self.scale_weight267        scale_input = self.scale_input268        if scale_weight is None:269            scale_weight = torch.ones((), device=input.device, dtype=torch.float32)270        else:271            scale_weight = scale_weight.to(input.device)272 273        if scale_input is None:274            scale_input = torch.ones((), device=input.device, dtype=torch.float32)275            input = torch.clamp(input, min=-448, max=448, out=input)276            input = input.reshape(-1, input_shape[2]).to(dtype)277        else:278            scale_input = scale_input.to(input.device)279            input = (input * (1.0 / scale_input).to(input_dtype)).reshape(-1, input_shape[2]).to(dtype)280 281        if bias is not None:282            o = torch._scaled_mm(input, w, out_dtype=input_dtype, bias=bias, scale_a=scale_input, scale_b=scale_weight)283        else:284            o = torch._scaled_mm(input, w, out_dtype=input_dtype, scale_a=scale_input, scale_b=scale_weight)285 286        if isinstance(o, tuple):287            o = o[0]288 289        if tensor_2d:290            return o.reshape(input_shape[0], -1)291 292        return o.reshape((-1, input_shape[1], self.weight.shape[0]))293 294    return None295 296class fp8_ops(manual_cast):297    class Linear(manual_cast.Linear):298        def reset_parameters(self):299            self.scale_weight = None300            self.scale_input = None301            return None302 303        def forward_comfy_cast_weights(self, input):304            out = fp8_linear(self, input)305            if out is not None:306                return out307 308            weight, bias = cast_bias_weight(self, input)309            return torch.nn.functional.linear(input, weight, bias)310 311def scaled_fp8_ops(fp8_matrix_mult=False, scale_input=False, override_dtype=None):312    logging.info("Using scaled fp8: fp8 matrix mult: {}, scale input: {}".format(fp8_matrix_mult, scale_input))313    class scaled_fp8_op(manual_cast):314        class Linear(manual_cast.Linear):315            def __init__(self, *args, **kwargs):316                if override_dtype is not None:317                    kwargs['dtype'] = override_dtype318                super().__init__(*args, **kwargs)319 320            def reset_parameters(self):321                if not hasattr(self, 'scale_weight'):322                    self.scale_weight = torch.nn.parameter.Parameter(data=torch.ones((), device=self.weight.device, dtype=torch.float32), requires_grad=False)323 324                if not scale_input:325                    self.scale_input = None326 327                if not hasattr(self, 'scale_input'):328                    self.scale_input = torch.nn.parameter.Parameter(data=torch.ones((), device=self.weight.device, dtype=torch.float32), requires_grad=False)329                return None330 331            def forward_comfy_cast_weights(self, input):332                if fp8_matrix_mult:333                    out = fp8_linear(self, input)334                    if out is not None:335                        return out336 337                weight, bias = cast_bias_weight(self, input)338 339                if weight.numel() < input.numel(): #TODO: optimize340                    return torch.nn.functional.linear(input, weight * self.scale_weight.to(device=weight.device, dtype=weight.dtype), bias)341                else:342                    return torch.nn.functional.linear(input * self.scale_weight.to(device=weight.device, dtype=weight.dtype), weight, bias)343 344            def convert_weight(self, weight, inplace=False, **kwargs):345                if inplace:346                    weight *= self.scale_weight.to(device=weight.device, dtype=weight.dtype)347                    return weight348                else:349                    return weight * self.scale_weight.to(device=weight.device, dtype=weight.dtype)350 351            def set_weight(self, weight, inplace_update=False, seed=None, **kwargs):352                weight = comfy.float.stochastic_rounding(weight / self.scale_weight.to(device=weight.device, dtype=weight.dtype), self.weight.dtype, seed=seed)353                if inplace_update:354                    self.weight.data.copy_(weight)355                else:356                    self.weight = torch.nn.Parameter(weight, requires_grad=False)357 358    return scaled_fp8_op359 360def pick_operations(weight_dtype, compute_dtype, load_device=None, disable_fast_fp8=False, fp8_optimizations=False, scaled_fp8=None):361    fp8_compute = comfy.model_management.supports_fp8_compute(load_device)362    if scaled_fp8 is not None:363        return scaled_fp8_ops(fp8_matrix_mult=fp8_compute and fp8_optimizations, scale_input=fp8_optimizations, override_dtype=scaled_fp8)364 365    if (366        fp8_compute and367        (fp8_optimizations or PerformanceFeature.Fp8MatrixMultiplication in args.fast) and368        not disable_fast_fp8369    ):370        return fp8_ops371 372    if compute_dtype is None or weight_dtype == compute_dtype:373        return disable_weight_init374 375    return manual_cast376