mirror of
https://github.com/lllyasviel/stable-diffusion-webui-forge.git
synced 2026-07-21 21:01:24 +08:00
71 lines
2.8 KiB
Python
71 lines
2.8 KiB
Python
import math
|
|
|
|
import torch
|
|
|
|
from backend import args, memory_management
|
|
from backend.modules.k_prediction import k_prediction_from_diffusers_scheduler
|
|
|
|
|
|
class KModel(torch.nn.Module):
|
|
def __init__(self, model: torch.nn.Module, diffusers_scheduler, k_predictor=None, config=None):
|
|
super().__init__()
|
|
|
|
self.config = config
|
|
|
|
self.storage_dtype = model.storage_dtype
|
|
self.computation_dtype = model.computation_dtype
|
|
|
|
_store = f"storage: {self.storage_dtype}"
|
|
_compute = f"computation: {self.computation_dtype}"
|
|
|
|
if args.dynamic_args.ops.endswith("Int8"):
|
|
_compute += f" + {torch.int8}"
|
|
if args.dynamic_args.ops.endswith("FP8"):
|
|
_compute += f" + {torch.float8_e4m3fn}"
|
|
if args.dynamic_args.ops.startswith("Mixed"):
|
|
_compute = "computation: Mixed"
|
|
|
|
memory_management.logger.info(f"Diffusion Model: {{{_store}, {_compute}}}")
|
|
|
|
self.diffusion_model = model
|
|
self.diffusion_model.eval()
|
|
self.diffusion_model.requires_grad_(False)
|
|
|
|
if k_predictor is None:
|
|
self.predictor = k_prediction_from_diffusers_scheduler(diffusers_scheduler)
|
|
else:
|
|
self.predictor = k_predictor
|
|
|
|
def apply_model(self, x, t, c_concat=None, c_crossattn=None, control=None, transformer_options={}, **kwargs):
|
|
sigma = t
|
|
xc = self.predictor.calculate_input(sigma, x)
|
|
if c_concat is not None:
|
|
xc = torch.cat([xc] + [c_concat], dim=1)
|
|
|
|
context = c_crossattn
|
|
dtype = self.computation_dtype
|
|
|
|
xc = xc.to(dtype)
|
|
t = self.predictor.timestep(t).float()
|
|
context = context.to(dtype)
|
|
extra_conds = {}
|
|
for o in kwargs:
|
|
extra = kwargs[o]
|
|
if hasattr(extra, "dtype"):
|
|
if extra.dtype != torch.int and extra.dtype != torch.long:
|
|
extra = extra.to(dtype)
|
|
extra_conds[o] = extra
|
|
|
|
model_output = self.diffusion_model(xc, t, context=context, control=control, transformer_options=transformer_options, **extra_conds).float()
|
|
return self.predictor.calculate_denoised(sigma, model_output, x)
|
|
|
|
def memory_required(self, input_shape: list[int]) -> float:
|
|
"""https://github.com/comfyanonymous/ComfyUI/blob/v0.3.64/comfy/model_base.py#L354"""
|
|
input_shapes = [input_shape]
|
|
area = sum(map(lambda input_shape: input_shape[0] * math.prod(input_shape[2:]), input_shapes))
|
|
|
|
if memory_management.xformers_enabled() or memory_management.pytorch_attention_flash_attention():
|
|
return (area * memory_management.dtype_size(self.computation_dtype) * 0.01 * self.config.memory_usage_factor) * (1024 * 1024)
|
|
else:
|
|
return (area * 0.15 * self.config.memory_usage_factor) * (1024 * 1024)
|