diff --git a/.github/workflows/pullrequest-ci-run.yml b/.github/workflows/pullrequest-ci-run.yml
index fd1ba16e2..691480bcf 100644
--- a/.github/workflows/pullrequest-ci-run.yml
+++ b/.github/workflows/pullrequest-ci-run.yml
@@ -35,3 +35,19 @@ jobs:
torch_version: ${{ matrix.torch_version }}
google_credentials: ${{ secrets.GCS_SERVICE_ACCOUNT_JSON }}
comfyui_flags: ${{ matrix.flags }}
+ use_prior_commit: 'true'
+ comment:
+ if: ${{ github.event.label.name == 'Run-CI-Test' }}
+ runs-on: ubuntu-latest
+ permissions:
+ pull-requests: write
+ steps:
+ - uses: actions/github-script@v6
+ with:
+ script: |
+ github.rest.issues.createComment({
+ issue_number: context.issue.number,
+ owner: context.repo.owner,
+ repo: context.repo.repo,
+ body: '(Automated Bot Message) CI Tests are running, you can view the results at https://ci.comfy.org/?branch=${{ github.event.pull_request.number }}%2Fmerge'
+ })
diff --git a/.gitignore b/.gitignore
index 5092c98f4..b212b0297 100644
--- a/.gitignore
+++ b/.gitignore
@@ -18,4 +18,5 @@ venv/
/tests-ui/data/object_info.json
/user/
*.log
-web_custom_versions/
\ No newline at end of file
+web_custom_versions/
+.DS_Store
diff --git a/comfy/controlnet.py b/comfy/controlnet.py
index 354b00367..dcfe492ce 100644
--- a/comfy/controlnet.py
+++ b/comfy/controlnet.py
@@ -34,6 +34,8 @@ import comfy.t2i_adapter.adapter
import comfy.ldm.cascade.controlnet
import comfy.cldm.mmdit
import comfy.ldm.hydit.controlnet
+import comfy.ldm.flux.controlnet_xlabs
+
def broadcast_image_to(tensor, target_batch_size, batched_number):
current_batch_size = tensor.shape[0]
@@ -416,6 +418,7 @@ def load_controlnet_mmdit(sd):
control = ControlNet(control_model, compression_ratio=1, latent_format=latent_format, load_device=load_device, manual_cast_dtype=manual_cast_dtype)
return control
+
def load_controlnet_hunyuandit(controlnet_data):
model_config, operations, load_device, unet_dtype, manual_cast_dtype = controlnet_config(controlnet_data)
@@ -427,6 +430,15 @@ def load_controlnet_hunyuandit(controlnet_data):
control = ControlNet(control_model, compression_ratio=1, latent_format=latent_format, load_device=load_device, manual_cast_dtype=manual_cast_dtype, extra_conds=extra_conds, strength_type=StrengthType.CONSTANT)
return control
+def load_controlnet_flux_xlabs(sd):
+ model_config, operations, load_device, unet_dtype, manual_cast_dtype = controlnet_config(sd)
+ control_model = comfy.ldm.flux.controlnet_xlabs.ControlNetFlux(operations=operations, device=load_device, dtype=unet_dtype, **model_config.unet_config)
+ control_model = controlnet_load_state_dict(control_model, sd)
+ extra_conds = ['y', 'guidance']
+ control = ControlNet(control_model, load_device=load_device, manual_cast_dtype=manual_cast_dtype, extra_conds=extra_conds)
+ return control
+
+
def load_controlnet(ckpt_path, model=None):
controlnet_data = comfy.utils.load_torch_file(ckpt_path, safe_load=True)
if 'after_proj_list.18.bias' in controlnet_data.keys(): #Hunyuan DiT
@@ -489,7 +501,10 @@ def load_controlnet(ckpt_path, model=None):
logging.warning("leftover keys: {}".format(leftover_keys))
controlnet_data = new_sd
elif "controlnet_blocks.0.weight" in controlnet_data: #SD3 diffusers format
- return load_controlnet_mmdit(controlnet_data)
+ if "double_blocks.0.img_attn.norm.key_norm.scale" in controlnet_data:
+ return load_controlnet_flux_xlabs(controlnet_data)
+ else:
+ return load_controlnet_mmdit(controlnet_data)
pth_key = 'control_model.zero_convs.0.0.weight'
pth = False
diff --git a/comfy/ldm/flux/controlnet_xlabs.py b/comfy/ldm/flux/controlnet_xlabs.py
new file mode 100644
index 000000000..3f40021b2
--- /dev/null
+++ b/comfy/ldm/flux/controlnet_xlabs.py
@@ -0,0 +1,104 @@
+#Original code can be found on: https://github.com/XLabs-AI/x-flux/blob/main/src/flux/controlnet.py
+
+import torch
+from torch import Tensor, nn
+from einops import rearrange, repeat
+
+from .layers import (DoubleStreamBlock, EmbedND, LastLayer,
+ MLPEmbedder, SingleStreamBlock,
+ timestep_embedding)
+
+from .model import Flux
+import comfy.ldm.common_dit
+
+
+class ControlNetFlux(Flux):
+ def __init__(self, image_model=None, dtype=None, device=None, operations=None, **kwargs):
+ super().__init__(final_layer=False, dtype=dtype, device=device, operations=operations, **kwargs)
+
+ # add ControlNet blocks
+ self.controlnet_blocks = nn.ModuleList([])
+ for _ in range(self.params.depth):
+ controlnet_block = operations.Linear(self.hidden_size, self.hidden_size, dtype=dtype, device=device)
+ # controlnet_block = zero_module(controlnet_block)
+ self.controlnet_blocks.append(controlnet_block)
+ self.pos_embed_input = operations.Linear(self.in_channels, self.hidden_size, bias=True, dtype=dtype, device=device)
+ self.gradient_checkpointing = False
+ self.input_hint_block = nn.Sequential(
+ operations.Conv2d(3, 16, 3, padding=1, dtype=dtype, device=device),
+ nn.SiLU(),
+ operations.Conv2d(16, 16, 3, padding=1, dtype=dtype, device=device),
+ nn.SiLU(),
+ operations.Conv2d(16, 16, 3, padding=1, stride=2, dtype=dtype, device=device),
+ nn.SiLU(),
+ operations.Conv2d(16, 16, 3, padding=1, dtype=dtype, device=device),
+ nn.SiLU(),
+ operations.Conv2d(16, 16, 3, padding=1, stride=2, dtype=dtype, device=device),
+ nn.SiLU(),
+ operations.Conv2d(16, 16, 3, padding=1, dtype=dtype, device=device),
+ nn.SiLU(),
+ operations.Conv2d(16, 16, 3, padding=1, stride=2, dtype=dtype, device=device),
+ nn.SiLU(),
+ operations.Conv2d(16, 16, 3, padding=1, dtype=dtype, device=device)
+ )
+
+ def forward_orig(
+ self,
+ img: Tensor,
+ img_ids: Tensor,
+ controlnet_cond: Tensor,
+ txt: Tensor,
+ txt_ids: Tensor,
+ timesteps: Tensor,
+ y: Tensor,
+ guidance: Tensor = None,
+ ) -> Tensor:
+ if img.ndim != 3 or txt.ndim != 3:
+ raise ValueError("Input img and txt tensors must have 3 dimensions.")
+
+ # running on sequences img
+ img = self.img_in(img)
+ controlnet_cond = self.input_hint_block(controlnet_cond)
+ controlnet_cond = rearrange(controlnet_cond, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2)
+ controlnet_cond = self.pos_embed_input(controlnet_cond)
+ img = img + controlnet_cond
+ vec = self.time_in(timestep_embedding(timesteps, 256))
+ if self.params.guidance_embed:
+ vec = vec + self.guidance_in(timestep_embedding(guidance, 256))
+ vec = vec + self.vector_in(y)
+ txt = self.txt_in(txt)
+
+ ids = torch.cat((txt_ids, img_ids), dim=1)
+ pe = self.pe_embedder(ids)
+
+ block_res_samples = ()
+
+ for block in self.double_blocks:
+ img, txt = block(img=img, txt=txt, vec=vec, pe=pe)
+ block_res_samples = block_res_samples + (img,)
+
+ controlnet_block_res_samples = ()
+ for block_res_sample, controlnet_block in zip(block_res_samples, self.controlnet_blocks):
+ block_res_sample = controlnet_block(block_res_sample)
+ controlnet_block_res_samples = controlnet_block_res_samples + (block_res_sample,)
+
+ return {"output": (controlnet_block_res_samples * 10)[:19]}
+
+ def forward(self, x, timesteps, context, y, guidance=None, hint=None, **kwargs):
+ hint = hint * 2.0 - 1.0
+
+ bs, c, h, w = x.shape
+ patch_size = 2
+ x = comfy.ldm.common_dit.pad_to_patch_size(x, (patch_size, patch_size))
+
+ img = rearrange(x, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=patch_size, pw=patch_size)
+
+ h_len = ((h + (patch_size // 2)) // patch_size)
+ w_len = ((w + (patch_size // 2)) // patch_size)
+ img_ids = torch.zeros((h_len, w_len, 3), device=x.device, dtype=x.dtype)
+ img_ids[..., 1] = img_ids[..., 1] + torch.linspace(0, h_len - 1, steps=h_len, device=x.device, dtype=x.dtype)[:, None]
+ img_ids[..., 2] = img_ids[..., 2] + torch.linspace(0, w_len - 1, steps=w_len, device=x.device, dtype=x.dtype)[None, :]
+ img_ids = repeat(img_ids, "h w c -> b (h w) c", b=bs)
+
+ txt_ids = torch.zeros((bs, context.shape[1], 3), device=x.device, dtype=x.dtype)
+ return self.forward_orig(img, img_ids, hint, context, txt_ids, timesteps, y, guidance)
diff --git a/comfy/ldm/flux/model.py b/comfy/ldm/flux/model.py
index db6cf3d22..b5373540a 100644
--- a/comfy/ldm/flux/model.py
+++ b/comfy/ldm/flux/model.py
@@ -38,7 +38,7 @@ class Flux(nn.Module):
Transformer model for flow matching on sequences.
"""
- def __init__(self, image_model=None, dtype=None, device=None, operations=None, **kwargs):
+ def __init__(self, image_model=None, final_layer=True, dtype=None, device=None, operations=None, **kwargs):
super().__init__()
self.dtype = dtype
params = FluxParams(**kwargs)
@@ -83,7 +83,8 @@ class Flux(nn.Module):
]
)
- self.final_layer = LastLayer(self.hidden_size, 1, self.out_channels, dtype=dtype, device=device, operations=operations)
+ if final_layer:
+ self.final_layer = LastLayer(self.hidden_size, 1, self.out_channels, dtype=dtype, device=device, operations=operations)
def forward_orig(
self,
@@ -94,6 +95,7 @@ class Flux(nn.Module):
timesteps: Tensor,
y: Tensor,
guidance: Tensor = None,
+ control=None,
) -> Tensor:
if img.ndim != 3 or txt.ndim != 3:
raise ValueError("Input img and txt tensors must have 3 dimensions.")
@@ -112,8 +114,15 @@ class Flux(nn.Module):
ids = torch.cat((txt_ids, img_ids), dim=1)
pe = self.pe_embedder(ids)
- for block in self.double_blocks:
- img, txt = block(img=img, txt=txt, vec=vec, pe=pe)
+ for i in range(len(self.double_blocks)):
+ img, txt = self.double_blocks[i](img=img, txt=txt, vec=vec, pe=pe)
+
+ if control is not None: #Controlnet
+ control_o = control.get("output")
+ if i < len(control_o):
+ add = control_o[i]
+ if add is not None:
+ img += add
img = torch.cat((txt, img), 1)
for block in self.single_blocks:
@@ -123,7 +132,7 @@ class Flux(nn.Module):
img = self.final_layer(img, vec) # (N, T, patch_size ** 2 * out_channels)
return img
- def forward(self, x, timestep, context, y, guidance, **kwargs):
+ def forward(self, x, timestep, context, y, guidance, control=None, **kwargs):
bs, c, h, w = x.shape
patch_size = 2
x = comfy.ldm.common_dit.pad_to_patch_size(x, (patch_size, patch_size))
@@ -138,5 +147,5 @@ class Flux(nn.Module):
img_ids = repeat(img_ids, "h w c -> b (h w) c", b=bs)
txt_ids = torch.zeros((bs, context.shape[1], 3), device=x.device, dtype=x.dtype)
- out = self.forward_orig(img, img_ids, context, txt_ids, timestep, y, guidance)
+ out = self.forward_orig(img, img_ids, context, txt_ids, timestep, y, guidance, control)
return rearrange(out, "b (h w) (c ph pw) -> b c (h ph) (w pw)", h=h_len, w=w_len, ph=2, pw=2)[:,:,:h,:w]
diff --git a/comfy/model_base.py b/comfy/model_base.py
index cb6949649..830bcc68c 100644
--- a/comfy/model_base.py
+++ b/comfy/model_base.py
@@ -1,3 +1,21 @@
+"""
+ This file is part of ComfyUI.
+ Copyright (C) 2024 Comfy
+
+ This program is free software: you can redistribute it and/or modify
+ it under the terms of the GNU General Public License as published by
+ the Free Software Foundation, either version 3 of the License, or
+ (at your option) any later version.
+
+ This program is distributed in the hope that it will be useful,
+ but WITHOUT ANY WARRANTY; without even the implied warranty of
+ MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+ GNU General Public License for more details.
+
+ You should have received a copy of the GNU General Public License
+ along with this program. If not, see .
+"""
+
import torch
import logging
from comfy.ldm.modules.diffusionmodules.openaimodel import UNetModel, Timestep
@@ -77,10 +95,13 @@ class BaseModel(torch.nn.Module):
self.device = device
if not unet_config.get("disable_unet_model_creation", False):
- if self.manual_cast_dtype is not None:
- operations = comfy.ops.manual_cast
+ if model_config.custom_operations is None:
+ if self.manual_cast_dtype is not None:
+ operations = comfy.ops.manual_cast
+ else:
+ operations = comfy.ops.disable_weight_init
else:
- operations = comfy.ops.disable_weight_init
+ operations = model_config.custom_operations
self.diffusion_model = unet_model(**unet_config, device=device, operations=operations)
if comfy.model_management.force_channels_last():
self.diffusion_model.to(memory_format=torch.channels_last)
diff --git a/comfy/model_detection.py b/comfy/model_detection.py
index 15e6b735c..c05975cc9 100644
--- a/comfy/model_detection.py
+++ b/comfy/model_detection.py
@@ -495,7 +495,12 @@ def model_config_from_diffusers_unet(state_dict):
def convert_diffusers_mmdit(state_dict, output_prefix=""):
out_sd = {}
- if 'transformer_blocks.0.attn.add_q_proj.weight' in state_dict: #SD3
+ if 'transformer_blocks.0.attn.norm_added_k.weight' in state_dict: #Flux
+ depth = count_blocks(state_dict, 'transformer_blocks.{}.')
+ depth_single_blocks = count_blocks(state_dict, 'single_transformer_blocks.{}.')
+ hidden_size = state_dict["x_embedder.bias"].shape[0]
+ sd_map = comfy.utils.flux_to_diffusers({"depth": depth, "depth_single_blocks": depth_single_blocks, "hidden_size": hidden_size}, output_prefix=output_prefix)
+ elif 'transformer_blocks.0.attn.add_q_proj.weight' in state_dict: #SD3
num_blocks = count_blocks(state_dict, 'transformer_blocks.{}.')
depth = state_dict["pos_embed.proj.weight"].shape[0] // 64
sd_map = comfy.utils.mmdit_to_diffusers({"depth": depth, "num_blocks": num_blocks}, output_prefix=output_prefix)
@@ -521,7 +526,12 @@ def convert_diffusers_mmdit(state_dict, output_prefix=""):
old_weight = out_sd.get(t[0], None)
if old_weight is None:
old_weight = torch.empty_like(weight)
- old_weight = old_weight.repeat([3] + [1] * (len(old_weight.shape) - 1))
+ if old_weight.shape[offset[0]] < offset[1] + offset[2]:
+ exp = list(weight.shape)
+ exp[offset[0]] = offset[1] + offset[2]
+ new = torch.empty(exp, device=weight.device, dtype=weight.dtype)
+ new[:old_weight.shape[0]] = old_weight
+ old_weight = new
w = old_weight.narrow(offset[0], offset[1], offset[2])
else:
diff --git a/comfy/model_management.py b/comfy/model_management.py
index a0105131d..fdf3308ba 100644
--- a/comfy/model_management.py
+++ b/comfy/model_management.py
@@ -438,11 +438,11 @@ def load_models_gpu(models, memory_required=0, force_patch_weights=False, minimu
global vram_state
inference_memory = minimum_inference_memory()
- extra_mem = max(inference_memory, memory_required) + 100 * 1024 * 1024
+ extra_mem = max(inference_memory, memory_required + 300 * 1024 * 1024)
if minimum_memory_required is None:
minimum_memory_required = extra_mem
else:
- minimum_memory_required = max(inference_memory, minimum_memory_required) + 100 * 1024 * 1024
+ minimum_memory_required = max(inference_memory, minimum_memory_required + 300 * 1024 * 1024)
models = set(models)
@@ -684,6 +684,20 @@ def text_encoder_device():
else:
return torch.device("cpu")
+def text_encoder_initial_device(load_device, offload_device, model_size=0):
+ if load_device == offload_device or model_size <= 1024 * 1024 * 1024:
+ return offload_device
+
+ if is_device_mps(load_device):
+ return offload_device
+
+ mem_l = get_free_memory(load_device)
+ mem_o = get_free_memory(offload_device)
+ if mem_l > (mem_o * 0.5) and model_size * 1.2 < mem_l:
+ return load_device
+ else:
+ return offload_device
+
def text_encoder_dtype(device=None):
if args.fp8_e4m3fn_text_enc:
return torch.float8_e4m3fn
diff --git a/comfy/model_patcher.py b/comfy/model_patcher.py
index ae3d20514..39d00d82e 100644
--- a/comfy/model_patcher.py
+++ b/comfy/model_patcher.py
@@ -355,13 +355,14 @@ class ModelPatcher:
return self.model
- def lowvram_load(self, device_to=None, lowvram_model_memory=0, force_patch_weights=False):
+ def lowvram_load(self, device_to=None, lowvram_model_memory=0, force_patch_weights=False, full_load=False):
mem_counter = 0
patch_counter = 0
lowvram_counter = 0
for n, m in self.model.named_modules():
lowvram_weight = False
- if hasattr(m, "comfy_cast_weights"):
+
+ if not full_load and hasattr(m, "comfy_cast_weights"):
module_mem = comfy.model_management.module_size(m)
if mem_counter + module_mem >= lowvram_model_memory:
lowvram_weight = True
@@ -401,13 +402,16 @@ class ModelPatcher:
if weight.device == device_to:
continue
- self.patch_weight_to_device(weight_key) #TODO: speed this up without OOM
- self.patch_weight_to_device(bias_key)
+ weight_to = None
+ if full_load:#TODO
+ weight_to = device_to
+ self.patch_weight_to_device(weight_key, device_to=weight_to) #TODO: speed this up without OOM
+ self.patch_weight_to_device(bias_key, device_to=weight_to)
m.to(device_to)
logging.debug("lowvram: loaded module regularly {} {}".format(n, m))
if lowvram_counter > 0:
- logging.info("loaded in lowvram mode {}".format(lowvram_model_memory / (1024 * 1024)))
+ logging.info("loaded partially {} {}".format(lowvram_model_memory / (1024 * 1024), patch_counter))
self.model.model_lowvram = True
else:
logging.info("loaded completely {} {}".format(lowvram_model_memory / (1024 * 1024), mem_counter / (1024 * 1024)))
@@ -665,12 +669,15 @@ class ModelPatcher:
return memory_freed
def partially_load(self, device_to, extra_memory=0):
+ self.unpatch_model(unpatch_weights=False)
+ self.patch_model(patch_weights=False)
+ full_load = False
if self.model.model_lowvram == False:
return 0
if self.model.model_loaded_weight_memory + extra_memory > self.model_size():
- pass #TODO: Full load
+ full_load = True
current_used = self.model.model_loaded_weight_memory
- self.lowvram_load(device_to, lowvram_model_memory=current_used + extra_memory)
+ self.lowvram_load(device_to, lowvram_model_memory=current_used + extra_memory, full_load=full_load)
return self.model.model_loaded_weight_memory - current_used
def current_loaded_device(self):
diff --git a/comfy/sd.py b/comfy/sd.py
index 94fc4e590..c8a2f086c 100644
--- a/comfy/sd.py
+++ b/comfy/sd.py
@@ -62,7 +62,7 @@ def load_lora_for_models(model, clip, lora, strength_model, strength_clip):
class CLIP:
- def __init__(self, target=None, embedding_directory=None, no_init=False, tokenizer_data={}):
+ def __init__(self, target=None, embedding_directory=None, no_init=False, tokenizer_data={}, parameters=0):
if no_init:
return
params = target.params.copy()
@@ -71,20 +71,24 @@ class CLIP:
load_device = model_management.text_encoder_device()
offload_device = model_management.text_encoder_offload_device()
- params['device'] = offload_device
dtype = model_management.text_encoder_dtype(load_device)
params['dtype'] = dtype
-
+ params['device'] = model_management.text_encoder_initial_device(load_device, offload_device, parameters * model_management.dtype_size(dtype))
self.cond_stage_model = clip(**(params))
for dt in self.cond_stage_model.dtypes:
if not model_management.supports_cast(load_device, dt):
load_device = offload_device
+ if params['device'] != offload_device:
+ self.cond_stage_model.to(offload_device)
+ logging.warning("Had to shift TE back.")
self.tokenizer = tokenizer(embedding_directory=embedding_directory, tokenizer_data=tokenizer_data)
self.patcher = comfy.model_patcher.ModelPatcher(self.cond_stage_model, load_device=load_device, offload_device=offload_device)
+ if params['device'] == load_device:
+ model_management.load_model_gpu(self.patcher)
self.layer_idx = None
- logging.debug("CLIP model load device: {}, offload device: {}".format(load_device, offload_device))
+ logging.debug("CLIP model load device: {}, offload device: {}, current: {}".format(load_device, offload_device, params['device']))
def clone(self):
n = CLIP(no_init=True)
@@ -456,7 +460,11 @@ def load_clip(ckpt_paths, embedding_directory=None, clip_type=CLIPType.STABLE_DI
clip_target.clip = comfy.text_encoders.sd3_clip.SD3ClipModel
clip_target.tokenizer = comfy.text_encoders.sd3_clip.SD3Tokenizer
- clip = CLIP(clip_target, embedding_directory=embedding_directory)
+ parameters = 0
+ for c in clip_data:
+ parameters += comfy.utils.calculate_parameters(c)
+
+ clip = CLIP(clip_target, embedding_directory=embedding_directory, parameters=parameters)
for c in clip_data:
m, u = clip.load_sd(c)
if len(m) > 0:
@@ -498,15 +506,19 @@ def load_checkpoint(config_path=None, ckpt_path=None, output_vae=True, output_cl
return (model, clip, vae)
-def load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, output_clipvision=False, embedding_directory=None, output_model=True):
+def load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, output_clipvision=False, embedding_directory=None, output_model=True, model_options={}):
sd = comfy.utils.load_torch_file(ckpt_path)
- sd_keys = sd.keys()
+ out = load_state_dict_guess_config(sd, output_vae, output_clip, output_clipvision, embedding_directory, output_model, model_options)
+ if out is None:
+ raise RuntimeError("ERROR: Could not detect model type of: {}".format(ckpt_path))
+ return out
+
+def load_state_dict_guess_config(sd, output_vae=True, output_clip=True, output_clipvision=False, embedding_directory=None, output_model=True, model_options={}):
clip = None
clipvision = None
vae = None
model = None
model_patcher = None
- clip_target = None
diffusion_model_prefix = model_detection.unet_prefix_from_state_dict(sd)
parameters = comfy.utils.calculate_parameters(sd, diffusion_model_prefix)
@@ -515,13 +527,18 @@ def load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, o
model_config = model_detection.model_config_from_unet(sd, diffusion_model_prefix)
if model_config is None:
- raise RuntimeError("ERROR: Could not detect model type of: {}".format(ckpt_path))
+ return None
unet_weight_dtype = list(model_config.supported_inference_dtypes)
if weight_dtype is not None:
unet_weight_dtype.append(weight_dtype)
- unet_dtype = model_management.unet_dtype(model_params=parameters, supported_dtypes=unet_weight_dtype)
+ model_config.custom_operations = model_options.get("custom_operations", None)
+ unet_dtype = model_options.get("weight_dtype", None)
+
+ if unet_dtype is None:
+ unet_dtype = model_management.unet_dtype(model_params=parameters, supported_dtypes=unet_weight_dtype)
+
manual_cast_dtype = model_management.unet_manual_cast(unet_dtype, load_device, model_config.supported_inference_dtypes)
model_config.set_inference_dtype(unet_dtype, manual_cast_dtype)
@@ -545,7 +562,8 @@ def load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, o
if clip_target is not None:
clip_sd = model_config.process_clip_state_dict(sd)
if len(clip_sd) > 0:
- clip = CLIP(clip_target, embedding_directory=embedding_directory, tokenizer_data=clip_sd)
+ parameters = comfy.utils.calculate_parameters(clip_sd)
+ clip = CLIP(clip_target, embedding_directory=embedding_directory, tokenizer_data=clip_sd, parameters=parameters)
m, u = clip.load_sd(clip_sd, full_model=True)
if len(m) > 0:
m_filter = list(filter(lambda a: ".logit_scale" not in a and ".transformer.text_projection.weight" not in a, m))
diff --git a/comfy/supported_models_base.py b/comfy/supported_models_base.py
index bc0a7e311..7a2152f91 100644
--- a/comfy/supported_models_base.py
+++ b/comfy/supported_models_base.py
@@ -1,3 +1,21 @@
+"""
+ This file is part of ComfyUI.
+ Copyright (C) 2024 Comfy
+
+ This program is free software: you can redistribute it and/or modify
+ it under the terms of the GNU General Public License as published by
+ the Free Software Foundation, either version 3 of the License, or
+ (at your option) any later version.
+
+ This program is distributed in the hope that it will be useful,
+ but WITHOUT ANY WARRANTY; without even the implied warranty of
+ MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+ GNU General Public License for more details.
+
+ You should have received a copy of the GNU General Public License
+ along with this program. If not, see .
+"""
+
import torch
from . import model_base
from . import utils
@@ -30,6 +48,7 @@ class BASE:
memory_usage_factor = 2.0
manual_cast_dtype = None
+ custom_operations = None
@classmethod
def matches(s, unet_config, state_dict=None):
diff --git a/comfy/utils.py b/comfy/utils.py
index e6736dbde..d0d410d97 100644
--- a/comfy/utils.py
+++ b/comfy/utils.py
@@ -457,8 +457,27 @@ def flux_to_diffusers(mmdit_config, output_prefix=""):
key_map["{}add_k_proj.{}".format(k, end)] = (qkv, (0, hidden_size, hidden_size))
key_map["{}add_v_proj.{}".format(k, end)] = (qkv, (0, hidden_size * 2, hidden_size))
- block_map = {"attn.to_out.0.weight": "img_attn.proj.weight",
- "attn.to_out.0.bias": "img_attn.proj.bias",
+ block_map = {
+ "attn.to_out.0.weight": "img_attn.proj.weight",
+ "attn.to_out.0.bias": "img_attn.proj.bias",
+ "norm1.linear.weight": "img_mod.lin.weight",
+ "norm1.linear.bias": "img_mod.lin.bias",
+ "norm1_context.linear.weight": "txt_mod.lin.weight",
+ "norm1_context.linear.bias": "txt_mod.lin.bias",
+ "attn.to_add_out.weight": "txt_attn.proj.weight",
+ "attn.to_add_out.bias": "txt_attn.proj.bias",
+ "ff.net.0.proj.weight": "img_mlp.0.weight",
+ "ff.net.0.proj.bias": "img_mlp.0.bias",
+ "ff.net.2.weight": "img_mlp.2.weight",
+ "ff.net.2.bias": "img_mlp.2.bias",
+ "ff_context.net.0.proj.weight": "txt_mlp.0.weight",
+ "ff_context.net.0.proj.bias": "txt_mlp.0.bias",
+ "ff_context.net.2.weight": "txt_mlp.2.weight",
+ "ff_context.net.2.bias": "txt_mlp.2.bias",
+ "attn.norm_q.weight": "img_attn.norm.query_norm.scale",
+ "attn.norm_k.weight": "img_attn.norm.key_norm.scale",
+ "attn.norm_added_q.weight": "txt_attn.norm.query_norm.scale",
+ "attn.norm_added_k.weight": "txt_attn.norm.key_norm.scale",
}
for k in block_map:
@@ -474,15 +493,41 @@ def flux_to_diffusers(mmdit_config, output_prefix=""):
key_map["{}to_q.{}".format(k, end)] = (qkv, (0, 0, hidden_size))
key_map["{}to_k.{}".format(k, end)] = (qkv, (0, hidden_size, hidden_size))
key_map["{}to_v.{}".format(k, end)] = (qkv, (0, hidden_size * 2, hidden_size))
- key_map["{}proj_mlp.{}".format(k, end)] = (qkv, (0, hidden_size * 3, hidden_size))
+ key_map["{}.proj_mlp.{}".format(prefix_from, end)] = (qkv, (0, hidden_size * 3, hidden_size * 4))
- block_map = {#TODO
+ block_map = {
+ "norm.linear.weight": "modulation.lin.weight",
+ "norm.linear.bias": "modulation.lin.bias",
+ "proj_out.weight": "linear2.weight",
+ "proj_out.bias": "linear2.bias",
+ "attn.norm_q.weight": "norm.query_norm.scale",
+ "attn.norm_k.weight": "norm.key_norm.scale",
}
for k in block_map:
key_map["{}.{}".format(prefix_from, k)] = "{}.{}".format(prefix_to, block_map[k])
- MAP_BASIC = { #TODO
+ MAP_BASIC = {
+ ("final_layer.linear.bias", "proj_out.bias"),
+ ("final_layer.linear.weight", "proj_out.weight"),
+ ("img_in.bias", "x_embedder.bias"),
+ ("img_in.weight", "x_embedder.weight"),
+ ("time_in.in_layer.bias", "time_text_embed.timestep_embedder.linear_1.bias"),
+ ("time_in.in_layer.weight", "time_text_embed.timestep_embedder.linear_1.weight"),
+ ("time_in.out_layer.bias", "time_text_embed.timestep_embedder.linear_2.bias"),
+ ("time_in.out_layer.weight", "time_text_embed.timestep_embedder.linear_2.weight"),
+ ("txt_in.bias", "context_embedder.bias"),
+ ("txt_in.weight", "context_embedder.weight"),
+ ("vector_in.in_layer.bias", "time_text_embed.text_embedder.linear_1.bias"),
+ ("vector_in.in_layer.weight", "time_text_embed.text_embedder.linear_1.weight"),
+ ("vector_in.out_layer.bias", "time_text_embed.text_embedder.linear_2.bias"),
+ ("vector_in.out_layer.weight", "time_text_embed.text_embedder.linear_2.weight"),
+ ("guidance_in.in_layer.bias", "time_text_embed.guidance_embedder.linear_1.bias"),
+ ("guidance_in.in_layer.weight", "time_text_embed.guidance_embedder.linear_1.weight"),
+ ("guidance_in.out_layer.bias", "time_text_embed.guidance_embedder.linear_2.bias"),
+ ("guidance_in.out_layer.weight", "time_text_embed.guidance_embedder.linear_2.weight"),
+ ("final_layer.adaLN_modulation.1.bias", "norm_out.linear.bias", swap_scale_shift),
+ ("final_layer.adaLN_modulation.1.weight", "norm_out.linear.weight", swap_scale_shift),
}
for k in MAP_BASIC: