227 lines
7.7 KiB
Python
227 lines
7.7 KiB
Python
import torch
|
|
import comfy
|
|
from functools import partial
|
|
from enum import Enum
|
|
from comfy.model_base import (
|
|
BaseModel,
|
|
convert_tensor,
|
|
QwenImage as _QwenImage,
|
|
Flux as _Flux,
|
|
Flux2 as _Flux2
|
|
)
|
|
from .piflow_policies import POLICY_CLASSES
|
|
from . import architectures
|
|
|
|
|
|
class ModelType(Enum):
|
|
PIFLOW = 1
|
|
|
|
|
|
class ModelSamplingPiFlow(torch.nn.Module):
|
|
def __init__(self, model_config=None):
|
|
super().__init__()
|
|
if model_config is not None:
|
|
sampling_settings = model_config.sampling_settings
|
|
else:
|
|
sampling_settings = {}
|
|
|
|
self.set_parameters(
|
|
shift=sampling_settings.get("shift", 3.2),
|
|
multiplier=sampling_settings.get("multiplier", 1.0),
|
|
patch_size=sampling_settings.get("patch_size", None)
|
|
)
|
|
|
|
def set_parameters(self, shift=3.2, multiplier=1.0, patch_size=None):
|
|
self.shift = shift
|
|
self.multiplier = multiplier
|
|
self.patch_size = patch_size
|
|
|
|
def timestep(self, sigma):
|
|
return sigma * self.multiplier
|
|
|
|
def warp_t(self, t):
|
|
shift = self.shift
|
|
return shift * t / (1 + (shift - 1) * t)
|
|
|
|
def unwarp_t(self, t):
|
|
shift = self.shift
|
|
return t / (shift + (1 - shift) * t)
|
|
|
|
def unpatchify(self, x):
|
|
if self.patch_size is None:
|
|
return x
|
|
elif len(self.patch_size) == 2:
|
|
assert x.dim() == 4, "Expected 4D tensor for 2D patchify"
|
|
b, c, h, w = x.shape
|
|
ph, pw = self.patch_size
|
|
x = x.reshape(
|
|
b, c // (ph * pw), ph, pw, h, w
|
|
).permute(0, 1, 4, 2, 5, 3).reshape(
|
|
b, c // (ph * pw), h * ph, w * pw)
|
|
return x
|
|
elif len(self.patch_size) == 3:
|
|
assert x.dim() == 5, "Expected 5D tensor for 3D patchify"
|
|
b, c, t, h, w = x.shape
|
|
pt, ph, pw = self.patch_size
|
|
x = x.reshape(
|
|
b, c // (pt * ph * pw), pt, ph, pw, t, h, w
|
|
).permute(0, 1, 5, 2, 6, 3, 7, 4).reshape(
|
|
b, c // (pt * ph * pw), t * pt, h * ph, w * pw)
|
|
return x
|
|
else:
|
|
raise ValueError("Unsupported patch size length")
|
|
|
|
def patchify(self, x):
|
|
if self.patch_size is None:
|
|
return x
|
|
elif len(self.patch_size) == 2:
|
|
assert x.dim() == 4, "Expected 4D tensor for 2D patchify"
|
|
b, c, h, w = x.shape
|
|
ph, pw = self.patch_size
|
|
x = x.reshape(
|
|
b, c, h // ph, ph, w // pw, pw
|
|
).permute(0, 1, 3, 5, 2, 4).reshape(
|
|
b, c * ph * pw, h // ph, w // pw)
|
|
return x
|
|
elif len(self.patch_size) == 3:
|
|
assert x.dim() == 5, "Expected 5D tensor for 3D patchify"
|
|
b, c, t, h, w = x.shape
|
|
pt, ph, pw = self.patch_size
|
|
x = x.reshape(
|
|
b, c, t // pt, pt, h // ph, ph, w // pw, pw
|
|
).permute(0, 1, 3, 5, 7, 2, 4, 6).reshape(
|
|
b, c * pt * ph * pw, t // pt, h // ph, w // pw)
|
|
return x
|
|
else:
|
|
raise ValueError("Unsupported patch size length")
|
|
|
|
def percent_to_sigma(self, percent):
|
|
if percent <= 0.0:
|
|
return 1.0
|
|
if percent >= 1.0:
|
|
return 0.0
|
|
return self.warp_t(1.0 - percent)
|
|
|
|
|
|
# stash original
|
|
_original_model_sampling = comfy.model_base.model_sampling
|
|
|
|
|
|
def model_sampling(model_config, model_type):
|
|
if model_type == ModelType.PIFLOW:
|
|
c = comfy.model_sampling.CONST
|
|
s = ModelSamplingPiFlow
|
|
|
|
class ModelSampling(s, c):
|
|
pass
|
|
|
|
return ModelSampling(model_config)
|
|
|
|
# fallback to original
|
|
return _original_model_sampling(model_config, model_type)
|
|
|
|
|
|
# patch comfyui model_sampling
|
|
comfy.model_base.model_sampling = model_sampling
|
|
|
|
|
|
class BasePiFlow(BaseModel):
|
|
|
|
def __init__(self, model_config, diffusion_model, model_type=ModelType.PIFLOW, device=None):
|
|
BaseModel.__init__(
|
|
self, model_config, model_type=model_type, device=device, unet_model=diffusion_model)
|
|
|
|
policy_config = model_config.policy_config.copy()
|
|
policy_type = policy_config.pop("type")
|
|
self.policy_class = partial(POLICY_CLASSES[policy_type], **policy_config)
|
|
|
|
def _apply_model(self, x, t, c_concat=None, c_crossattn=None, control=None, transformer_options={}, **kwargs):
|
|
sigma = t
|
|
xc = self.model_sampling.calculate_input(sigma, x)
|
|
|
|
if c_concat is not None:
|
|
xc = torch.cat([xc] + [comfy.model_management.cast_to_device(c_concat, xc.device, xc.dtype)], dim=1)
|
|
|
|
context = c_crossattn
|
|
dtype = self.get_dtype_inference()
|
|
|
|
xc = xc.to(dtype)
|
|
device = xc.device
|
|
t = self.model_sampling.timestep(t).float()
|
|
if context is not None:
|
|
context = comfy.model_management.cast_to_device(context, device, dtype)
|
|
|
|
extra_conds = {}
|
|
for o in kwargs:
|
|
extra = kwargs[o]
|
|
|
|
if hasattr(extra, "dtype"):
|
|
extra = convert_tensor(extra, dtype, device)
|
|
elif isinstance(extra, list):
|
|
ex = []
|
|
for ext in extra:
|
|
ex.append(convert_tensor(ext, dtype, device))
|
|
extra = ex
|
|
extra_conds[o] = extra
|
|
|
|
t = self.process_timestep(t, x=x, **extra_conds)
|
|
if "latent_shapes" in extra_conds:
|
|
raise RuntimeError("Packed/nested latents are not supported by pi-Flow models.")
|
|
|
|
transformer_options = transformer_options.copy()
|
|
transformer_options["prefetch_dynamic_vbars"] = (
|
|
self.current_patcher is not None and self.current_patcher.is_dynamic()
|
|
)
|
|
|
|
model_output = self.diffusion_model(xc, t, context=context, control=control,
|
|
transformer_options=transformer_options, **extra_conds)
|
|
if isinstance(model_output, dict):
|
|
model_output = {k: v.float() for k, v in model_output.items()}
|
|
else:
|
|
model_output = model_output.float()
|
|
return self.policy_class(model_output, x, sigma)
|
|
|
|
|
|
class GMQwenImage(_QwenImage, BasePiFlow):
|
|
|
|
def __init__(self, model_config, device=None):
|
|
BasePiFlow.__init__(self, model_config, architectures.GMQwenImageTransformer2DModel, device=device)
|
|
self.memory_usage_factor_conds = ("ref_latents",)
|
|
|
|
|
|
class QwenImage(_QwenImage, BasePiFlow):
|
|
|
|
def __init__(self, model_config, device=None):
|
|
BasePiFlow.__init__(self, model_config, architectures.QwenImageTransformer2DModelMod, device=device)
|
|
self.memory_usage_factor_conds = ("ref_latents",)
|
|
|
|
|
|
class GMFlux(_Flux, BasePiFlow):
|
|
def __init__(self, model_config, device=None):
|
|
BasePiFlow.__init__(self, model_config, architectures.GMFlux, device=device)
|
|
self.memory_usage_factor_conds = ("ref_latents",)
|
|
|
|
|
|
class Flux(_Flux, BasePiFlow):
|
|
def __init__(self, model_config, device=None):
|
|
BasePiFlow.__init__(self, model_config, architectures.FluxMod, device=device)
|
|
self.memory_usage_factor_conds = ("ref_latents",)
|
|
|
|
|
|
class GMFlux2(_Flux2, BasePiFlow):
|
|
def __init__(self, model_config, device=None):
|
|
BasePiFlow.__init__(self, model_config, architectures.GMFlux, device=device)
|
|
self.memory_usage_factor_conds = ("ref_latents",)
|
|
|
|
|
|
class Flux2(_Flux2, BasePiFlow):
|
|
def __init__(self, model_config, device=None):
|
|
BasePiFlow.__init__(self, model_config, architectures.FluxMod, device=device)
|
|
self.memory_usage_factor_conds = ("ref_latents",)
|
|
|
|
|
|
class AsymFlux2(_Flux2):
|
|
def __init__(self, model_config, device=None):
|
|
_Flux.__init__(self, model_config, device=device, unet_model=architectures.AsymFlux)
|
|
self.memory_usage_factor_conds = ("ref_latents",)
|