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: