xinyun/ComfyUI
1
0
Fork 0
mirror of https://git.datalinker.icu/comfyanonymous/ComfyUI synced 2026-10-01 13:57:10 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'comfyanonymous:master' into master

Browse Source
This commit is contained in:
Karol Stępień 2024-10-24 17:58:38 +02:00 committed by GitHub
commit 3f5eb772b5
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 26 additions and 21 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

21
comfy/controlnet.py
View File

@ -60,7 +60,7 @@ class StrengthType(Enum):
LINEAR_UP = 2 LINEAR_UP = 2
class ControlBase: class ControlBase:
def __init__(self, device=None): def __init__(self):
self.cond_hint_original = None self.cond_hint_original = None
self.cond_hint = None self.cond_hint = None
self.strength = 1.0 self.strength = 1.0
@ -72,10 +72,6 @@ class ControlBase:
self.compression_ratio = 8 self.compression_ratio = 8
self.upscale_algorithm = 'nearest-exact' self.upscale_algorithm = 'nearest-exact'
self.extra_args = {} self.extra_args = {}
if device is None:
device = comfy.model_management.get_torch_device()
self.device = device
self.previous_controlnet = None self.previous_controlnet = None
self.extra_conds = [] self.extra_conds = []
self.strength_type = StrengthType.CONSTANT self.strength_type = StrengthType.CONSTANT
@ -185,8 +181,8 @@ class ControlBase:
class ControlNet(ControlBase): class ControlNet(ControlBase):
def __init__(self, control_model=None, global_average_pooling=False, compression_ratio=8, latent_format=None, device=None, load_device=None, manual_cast_dtype=None, extra_conds=["y"], strength_type=StrengthType.CONSTANT, concat_mask=False): def __init__(self, control_model=None, global_average_pooling=False, compression_ratio=8, latent_format=None, load_device=None, manual_cast_dtype=None, extra_conds=["y"], strength_type=StrengthType.CONSTANT, concat_mask=False):
super().__init__(device) super().__init__()
self.control_model = control_model self.control_model = control_model
self.load_device = load_device self.load_device = load_device
if control_model is not None: if control_model is not None:
@ -242,7 +238,7 @@ class ControlNet(ControlBase):
to_concat.append(comfy.utils.repeat_to_batch_size(c, self.cond_hint.shape[0])) to_concat.append(comfy.utils.repeat_to_batch_size(c, self.cond_hint.shape[0]))
self.cond_hint = torch.cat([self.cond_hint] + to_concat, dim=1) self.cond_hint = torch.cat([self.cond_hint] + to_concat, dim=1)
self.cond_hint = self.cond_hint.to(device=self.device, dtype=dtype) self.cond_hint = self.cond_hint.to(device=x_noisy.device, dtype=dtype)
if x_noisy.shape[0] != self.cond_hint.shape[0]: if x_noisy.shape[0] != self.cond_hint.shape[0]:
self.cond_hint = broadcast_image_to(self.cond_hint, x_noisy.shape[0], batched_number) self.cond_hint = broadcast_image_to(self.cond_hint, x_noisy.shape[0], batched_number)
@ -341,8 +337,8 @@ class ControlLoraOps:
class ControlLora(ControlNet): class ControlLora(ControlNet):
def __init__(self, control_weights, global_average_pooling=False, device=None, model_options={}): #TODO? model_options def __init__(self, control_weights, global_average_pooling=False, model_options={}): #TODO? model_options
ControlBase.__init__(self, device) ControlBase.__init__(self)
self.control_weights = control_weights self.control_weights = control_weights
self.global_average_pooling = global_average_pooling self.global_average_pooling = global_average_pooling
self.extra_conds += ["y"] self.extra_conds += ["y"]
@ -662,12 +658,15 @@ def load_controlnet(ckpt_path, model=None, model_options={}):
class T2IAdapter(ControlBase): class T2IAdapter(ControlBase):
def __init__(self, t2i_model, channels_in, compression_ratio, upscale_algorithm, device=None): def __init__(self, t2i_model, channels_in, compression_ratio, upscale_algorithm, device=None):
super().__init__(device) super().__init__()
self.t2i_model = t2i_model self.t2i_model = t2i_model
self.channels_in = channels_in self.channels_in = channels_in
self.control_input = None self.control_input = None
self.compression_ratio = compression_ratio self.compression_ratio = compression_ratio
self.upscale_algorithm = upscale_algorithm self.upscale_algorithm = upscale_algorithm
if device is None:
device = comfy.model_management.get_torch_device()
self.device = device
def scale_image_to(self, width, height): def scale_image_to(self, width, height):
unshuffle_amount = self.t2i_model.unshuffle_amount unshuffle_amount = self.t2i_model.unshuffle_amount

18
comfy/ldm/modules/diffusionmodules/mmdit.py
View File

@ -5,7 +5,7 @@ from typing import Dict, Optional
import numpy as np import numpy as np
import torch import torch
import torch.nn as nn import torch.nn as nn
from .. import attention from ..attention import optimized_attention
from einops import rearrange, repeat from einops import rearrange, repeat
from .util import timestep_embedding from .util import timestep_embedding
import comfy.ops import comfy.ops
@ -266,8 +266,6 @@ def split_qkv(qkv, head_dim):
qkv = qkv.reshape(qkv.shape[0], qkv.shape[1], 3, -1, head_dim).movedim(2, 0) qkv = qkv.reshape(qkv.shape[0], qkv.shape[1], 3, -1, head_dim).movedim(2, 0)
return qkv[0], qkv[1], qkv[2] return qkv[0], qkv[1], qkv[2]
def optimized_attention(qkv, num_heads):
return attention.optimized_attention(qkv[0], qkv[1], qkv[2], num_heads)
class SelfAttention(nn.Module): class SelfAttention(nn.Module):
ATTENTION_MODES = ("xformers", "torch", "torch-hb", "math", "debug") ATTENTION_MODES = ("xformers", "torch", "torch-hb", "math", "debug")
@ -326,9 +324,9 @@ class SelfAttention(nn.Module):
return x return x
def forward(self, x: torch.Tensor) -> torch.Tensor: def forward(self, x: torch.Tensor) -> torch.Tensor:
qkv = self.pre_attention(x) q, k, v = self.pre_attention(x)
x = optimized_attention( x = optimized_attention(
qkv, num_heads=self.num_heads q, k, v, heads=self.num_heads
) )
x = self.post_attention(x) x = self.post_attention(x)
return x return x
@ -531,8 +529,8 @@ class DismantledBlock(nn.Module):
assert not self.pre_only assert not self.pre_only
qkv, intermediates = self.pre_attention(x, c) qkv, intermediates = self.pre_attention(x, c)
attn = optimized_attention( attn = optimized_attention(
qkv, qkv[0], qkv[1], qkv[2],
num_heads=self.attn.num_heads, heads=self.attn.num_heads,
) )
return self.post_attention(attn, *intermediates) return self.post_attention(attn, *intermediates)
@ -557,8 +555,8 @@ def _block_mixing(context, x, context_block, x_block, c):
qkv = tuple(o) qkv = tuple(o)
attn = optimized_attention( attn = optimized_attention(
qkv, qkv[0], qkv[1], qkv[2],
num_heads=x_block.attn.num_heads, heads=x_block.attn.num_heads,
) )
context_attn, x_attn = ( context_attn, x_attn = (
attn[:, : context_qkv[0].shape[1]], attn[:, : context_qkv[0].shape[1]],
@ -642,7 +640,7 @@ class SelfAttentionContext(nn.Module):
def forward(self, x): def forward(self, x):
qkv = self.qkv(x) qkv = self.qkv(x)
q, k, v = split_qkv(qkv, self.dim_head) q, k, v = split_qkv(qkv, self.dim_head)
x = optimized_attention((q.reshape(q.shape[0], q.shape[1], -1), k, v), self.heads) x = optimized_attention(q.reshape(q.shape[0], q.shape[1], -1), k, v, heads=self.heads)
return self.proj(x) return self.proj(x)
class ContextProcessorBlock(nn.Module): class ContextProcessorBlock(nn.Module):

4
comfy/lora.py
View File

@ -317,6 +317,10 @@ def model_lora_keys_unet(model, key_map={}):
key_lora = "lora_transformer_{}".format(k[:-len(".weight")].replace(".", "_")) #OneTrainer lora key_lora = "lora_transformer_{}".format(k[:-len(".weight")].replace(".", "_")) #OneTrainer lora
key_map[key_lora] = to key_map[key_lora] = to
key_lora = "lycoris_{}".format(k[:-len(".weight")].replace(".", "_")) #simpletuner lycoris format
key_map[key_lora] = to
if isinstance(model, comfy.model_base.AuraFlow): #Diffusers lora AuraFlow if isinstance(model, comfy.model_base.AuraFlow): #Diffusers lora AuraFlow
diffusers_keys = comfy.utils.auraflow_to_diffusers(model.model_config.unet_config, output_prefix="diffusion_model.") diffusers_keys = comfy.utils.auraflow_to_diffusers(model.model_config.unet_config, output_prefix="diffusion_model.")
for k in diffusers_keys: for k in diffusers_keys:

4
comfy/ops.py
View File

@ -264,10 +264,14 @@ def fp8_linear(self, input):
scale_input = self.scale_input scale_input = self.scale_input
if scale_weight is None: if scale_weight is None:
scale_weight = torch.ones((), device=input.device, dtype=torch.float32) scale_weight = torch.ones((), device=input.device, dtype=torch.float32)
else:
scale_weight = scale_weight.to(input.device)
if scale_input is None: if scale_input is None:
scale_input = torch.ones((), device=input.device, dtype=torch.float32) scale_input = torch.ones((), device=input.device, dtype=torch.float32)
inn = input.reshape(-1, input.shape[2]).to(dtype) inn = input.reshape(-1, input.shape[2]).to(dtype)
else: else:
scale_input = scale_input.to(input.device)
inn = (input * (1.0 / scale_input).to(input.dtype)).reshape(-1, input.shape[2]).to(dtype) inn = (input * (1.0 / scale_input).to(input.dtype)).reshape(-1, input.shape[2]).to(dtype)
if bias is not None: if bias is not None: