[
{
"path": ".idea/.gitignore",
"content": "# Default ignored files\n/shelf/\n/workspace.xml\n"
},
{
"path": ".idea/MistoControlNet-Flux-dev.iml",
"content": "\n\n \n \n \n \n \n"
},
{
"path": ".idea/inspectionProfiles/profiles_settings.xml",
"content": "\n \n \n \n \n"
},
{
"path": ".idea/misc.xml",
"content": "\n\n \n \n \n \n"
},
{
"path": ".idea/modules.xml",
"content": "\n\n \n \n \n \n \n"
},
{
"path": ".idea/vcs.xml",
"content": "\n\n \n \n \n"
},
{
"path": "README.md",
"content": "# New Update!\nWe have just launched our latest product, Misto.\nThe most powerful AI Mind Palace built for all designers.\nWarmly welcome everyone to try it out.\n\n### Website here: https://themisto.ai/\n\n \n\n[中文版-README](README_CN.md)\n## VERY IMPORTANT\n!!!Please update the ComfyUI-suite for fixed the tensor mismatch promblem. \n!!!please donot use AUTO cfg for our ksampler, it will have a very bad result. \n!!!Strength and prompt senstive, be care for your prompt and try 0.5 as the starting controlnet strength\n!!!update a new example workflow in workflow folder, get start with it. \n\n## Summary\nby TheMisto.ai @Shenzhen, China \nThis is a ControlNet network designed for any lineart or outline sketches, compatible with Flux1.dev. The ControlNet model parameters are approximately 1.4B. \n\nThis model is not compatible with XLabs loaders and samplers. Please use TheMisto.ai Flux ControlNet ComfyUI suite.\nThis is a Flow matching structure Flux-dev model, utilizing a scalable Transformer module as the backbone of this ControlNet. \n\nWe've implemented a dual-stream Transformer structure, which enhances alignment and expressiveness for various types of lineart and outline conditions without increasing inference time. The model has also been trained for alignment with both T5 and clip-l TextEncoders, ensuring balanced performance between conditioning images and text prompts. \nFor more details on the Flux.dev model structure, visit: https://huggingface.co/black-forest-labs/FLUX.1-dev \n\nThis ControlNet is compatible with Flux1.dev's fp16/fp8 and other models quantized with Flux1.dev. ByteDance 8/16-step distilled models have not been tested. \n- The example workflow uses the flux1-dev-Q4_K_S.gguf quantized model. \n- Generation quality: Flux1.dev(fp16)>>Flux1.dev(fp8)>>Other quantized models\n- Generation speed: Flux1.dev(fp16)<<< Flux1.dev(fp8) <<< Other quantized models\n\n## Performance\n### Performance Across Different Sizes and Scenarios\nTested in various common scenarios such as industrial design, architecture, interior design, animation, games, and photography. \nMake sure to craft your prompts well—precision is more important than length! \nPerformance examples are shown below: \n \n\n### Performance with Different Lineart or Scribble Preprocessors\nTest Parameters: \n- Prompt: \"Hyper-realistic 3D render of a classic Formula 1 race car, bright red with Marlboro and Agip logos, number 1, large black tires, dramatic studio lighting, dark moody background, reflective floor, cinematic atmosphere, Octane render style, high detail\"\n- controlnet_strength: 0.650.8 (Recommended: Anyline with 0.60.7)\n- steps: 30\n- guidance: 4.0\n- The quality of generated images is positively correlated with prompt quality. Controlnet_strength may vary for different types of lineart and outlines, so experiment with the settings!\n \n\n### Recommended Settings\n- Image resolution: 720px or above on the short edge\n- controlnet strength: 0.6~0.85 (adjust as needed)\n- guidance: 3.0~5.0 (adjust as needed)\n- steps: 30 or more\n\n## Huggingface (抱抱脸):\n[MistoLine_Flux.dev_v1](https://huggingface.co/TheMistoAI/MistoLine_Flux.dev)\n\n## Usage\n- Download the model from [MistoLine_Flux.dev_v1](https://huggingface.co/TheMistoAI/MistoLine_Flux.dev)\n- Place the model in the ComfyUI\\models\\TheMisto_model\\ directory\n- The directory will be automatically created the first time you run the ComfyUI's TheMisto.ai Flux ControlNet ComfyUI suite\n- Run using ComfyUI; an example workflow can be found in the workflow folder\n- Note: The length and width of the conditioning image must be divisible by 16, or an error will occur.\n### ComfyUI\n \n\n## Training Details\nThe Transformer structure, under the scale law, has a significant impact on training time and computational power (higher compute cost, longer training time). \nThe training cost for MistoLine_Flux1_dev is several times that of MistoLineSDXL.\nWe conducted extensive ablation experiments to balance performance with training costs. \nThis training was done using A100-80GB with bf16 mixed precision on the Flux1.dev series models. Apart from Lora, consumer-grade GPUs are basically unsuitable for training. \nIn our experiments with larger parameter models, multi-GPUs, multi-node parallel training was required, which is costly. \nIf we reach 50,000 stars, we will open-source the Technical Report detailing more training details.\n\n## License\nAlign to the FLUX.1 [dev] Non-Commercial License \nThis ComfyUI node falls under ComfyUI \nThis model is for research and educational purposes only and may not be used for any form of commercial purposes. \n\n## Business Cooperation\nFor any custom model training, commercial cooperation, AI application development, or other business collaboration matters, please contact us. \n\n- *Business:* info@themisto.ai\n\n## Media\n### International:\nWebsite: https://themisto.ai/ \nDiscord: https://discord.gg/fTyDB2CU \nX: https://x.com/AiThemisto79359 \n\n### Mainland China (中国大陆):\n*Website:* https://themisto.ai/ \n*WeChat Official Account:* TheMisto AI (Shenzhen Mixed Tuple Technology Co., Ltd.) \n*Xiaohongshu:* TheMisto.ai (Xiaohongshu ID: 4211745997) \n"
},
{
"path": "README_CN.md",
"content": " \n## 非常重要\n!!!请更新ComfyUI套件至最新版本，修复了tensor mismatch问题 \n!!!不要使用Auto CFG来跟Ksampler配合使用，可能会导致非常糟糕的结果 \n!!!对controlnet strength和prompt非常敏感，请好好写prompt和从0.5的controlnet strength逐步尝试 \n!!!更新了一个工作流在./workflow文件夹中，从从这个工作流开始 \n\n\n\n## 概述\nby TheMisto.ai @Shenzhen, China \n这是一个适用于任意线稿、轮廓用于Flux1.dev的ControlNet网络，本ControlNet参数约为1.4B。 \n本模型不兼容XLabs的加载和采样器，请使用TheMisto.ai Flux ControlNet ComfyUI套件。 \n\n这是一个Flow matching结构的Flux-dev模型，使用了可scale的Transformer 模块来作为本ControlNet的骨干网\n这次我们使用了双流型Transformer结构，在不增加推理时间的情况下对不同类型的线稿和轮廓条件有更好的表现力和对齐效果，同时对T5和clip-l两个TextEncoder的文字对齐也得到了对应的训练，不会出现只对conditioning image有响应而对prompt对齐能力下降的问题，尝试做到条件图像和文本对齐都有较好的表现。Flux.dev模型结构等请浏览：https://huggingface.co/black-forest-labs/FLUX.1-dev\n\n本ControlNet适用于Flux1.dev的fp16/fp8以及其他使用Flux1.dev量化的模型, 字节跳动8/16步蒸馏的没有测试过。 \n示例Workflow使用的是flux1-dev-Q4_K_S.gguf量化的模型。 \n生成质量：Flux1.dev(fp16)>>Flux1.dev(fp8)>>其他量化模型 \n生成速度：Flux1.dev(fp16)<<< Flux1.dev(fp8) <<< 其他量化模型\n\n### 效果\n#### 不同尺寸和场景效果\n测试了工业设计、建筑设计、室内设计、动漫、游戏、照片等常用场景。 \n请好好写prompt，不是长，是要比较精确！ \n效果如下： \n \n\n#### 不同类型的Lineart or scribble preprocessor\n测试参数:\n- Prompt: \"Hyper-realistic 3D render of a classic Formula 1 race car, bright red with Marlboro and Agip logos, number 1, large black tires, dramatic studio lighting, dark moody background, reflective floor, cinematic atmosphere, Octane render style, high detail\"\n- controlnet_strength: 0.65-0.8 (推荐：Anyline with 0.6-0.7)\n- step: 30\n- guidance: 4.0 \n- 生成效果跟prompt质量成正相关关系，不同类型的线稿、轮廓使用的controlnet_strength也可能不同，多试一下！\n \n\n\n### 推荐参数\n- 图像分辨率：720px 以上 短边 \n- controlnet strength：0.6~0.85 （视情况调整） \n- guidance:3.0~5.0 （视情况调整） \n- step:30 以上\n\n### Huggingface（抱抱脸）:\n[MistoLine_Flux.dev_v1](https://huggingface.co/TheMistoAI/MistoLine_Flux.dev)\n\n\n## 用法\n- 从[MistoLine_Flux.dev_v1](https://huggingface.co/TheMistoAI/MistoLine_Flux.dev)下载模型\n- 将模型放入到 ComfyUI\\models\\TheMisto_model\\ 目录下\n- 第一次运行ComfyUI的TheMisto.ai Flux ControlNet ComfyUI套件会自动创建该目录\n- 使用ComfyUI运行，文件夹workflow下有example的workflow \n- 注意：条件图的长和宽必须能被16整除，否则会报错\n### ComfyUI\n \n\n## 训练细节\nTransformer结构在scale law的背景下会对训练时间和算力产生巨大影响（更大的算力消耗，更多的训练时间），MistoLine_Flux1_dev的训练成本为MistoLineSDXL的数倍之多。我们做了大量的消融实验来确保效果与训练成本的平衡。 \n本次训练使用了A100-80GB，bf16混合精度，Flux1.dev系列模型的训练，训练的话除了Lora之外基本可以告别消费级显卡。\n在我们更大的参数量模型的实验中，需要使用多卡多节点并行训练，成本较大。 \n标星到5万我们将开源Technical report，介绍更详细的训练细节。\n\n## 许可\n- Align to the FLUX.1 [dev] Non-Commercial License \n- This ComfyUI node fall under ComfyUI \n- 本模型仅供研究和学习，不可用于任何形式商用\n\n## Business Cooperation（商务合作）\nFor any custom model training, commercial cooperation, AI application development, or other business collaboration matters. \nplease contact us. \n- Business :info@themisto.ai \n- Investment: investment@themisto.ai\n\n如果有任何模型定制训练，商业合作，AI应用开发等合作事宜请联系。\n同时我们也欢迎投资者咨询更多信息。 \n- 业务电邮：info@themisto.ai \n- 投资者关系：investment@themisto.ai\n\n## WIP\n- Flux1.dev-MistoCN-collection \n- Flux1.dev-Misto-IPAdapter \n\n你的star是我们开源的动力！\n\n## One more thing\n \n我们将会在最近推出我们自己的的产品：（一款极其简单易用的多模态AI创意创作APP - [Misto]） \n以最简单和启发性的体验，重新激发大众创作欲望 \n创意触手可及，扩展想象力边界，让无限灵感成就超级个体！\n\n支持平台：全平台 \n\n\n## 媒体\n### 国际化: \nwebsite: https://themisto.ai/ \nDiscord：https://discord.gg/fTyDB2CU \nX：https://x.com/AiThemisto79359\n\n### 中国大陆地区:\nwebsite: https://themisto.ai/ \nwechat公众号：TheMisto AI（深圳混合元组科技有限公司） \n小红书：TheMisto.ai (小红书号：4211745997)\n"
},
{
"path": "__init__.py",
"content": "from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS\n\n__all__ = [\"NODE_CLASS_MAPPINGS\", \"NODE_DISPLAY_NAME_MAPPINGS\"]"
},
{
"path": "modules/misto_controlnet.py",
"content": "import torch\nfrom diffusers.utils import is_torch_version\nfrom torch import Tensor, nn\nfrom einops import rearrange\nfrom typing import Any, Dict, Tuple, Union\nfrom .utils import EmbedND, MLPEmbedder, DoubleStreamBlock, SingleStreamBlock, timestep_embedding\n\ndef zero_module(module):\n for p in module.parameters():\n nn.init.zeros_(p)\n return module\n\n\nclass CondDownsamplBlock(nn.Module):\n def __init__(self):\n super().__init__()\n self.encoder = nn.Sequential(\n nn.Conv2d(3, 16, 3, padding=1),\n nn.SiLU(),\n nn.Conv2d(16, 16, 1),\n nn.SiLU(),\n nn.Conv2d(16, 16, 3, padding=1),\n nn.SiLU(),\n nn.Conv2d(16, 16, 3, padding=1, stride=2),\n nn.SiLU(),\n nn.Conv2d(16, 16, 3, padding=1),\n nn.SiLU(),\n nn.Conv2d(16, 16, 3, padding=1, stride=2),\n nn.SiLU(),\n nn.Conv2d(16, 16, 3, padding=1),\n nn.SiLU(),\n nn.Conv2d(16, 16, 3, padding=1, stride=2),\n nn.SiLU(),\n nn.Conv2d(16, 16, 1),\n nn.SiLU(),\n zero_module(nn.Conv2d(16, 16, 3, padding=1))\n )\n\n def forward(self, x):\n return self.encoder(x)\n\n\nclass EnhanceControlnet(nn.Module):\n def __init__(self, hidden_size):\n super().__init__()\n self.linear = nn.Linear(hidden_size, hidden_size)\n self.act = nn.SiLU()\n nn.init.eye_(self.linear.weight)\n nn.init.zeros_(self.linear.bias)\n\n def forward(self, x):\n return self.act(self.linear(x))\n\n\n\nclass MistoControlNetFluxDev(nn.Module):\n _supports_gradient_checkpointing = True\n\n def __init__(\n self,\n in_channels=64,\n vec_in_dim=768,\n context_in_dim=4096,\n hidden_size=3072,\n num_heads=24,\n num_transformer=2,\n num_single_transformer=2,\n guidance_embed=True,\n ):\n super().__init__()\n self.out_channels = in_channels\n self.axes_dim = [16, 56, 56]\n self.theta=10_000\n self.guidance_embed = guidance_embed\n\n if hidden_size % num_heads != 0:\n raise ValueError(f\"Hidden size {hidden_size} must be divisible by num_heads {num_heads}\")\n\n pe_dim = hidden_size // num_heads\n\n if sum(self.axes_dim) != pe_dim:\n raise ValueError(f\"Got {self.axes_dim} but expected positional dim {pe_dim}\")\n\n self.hidden_size = hidden_size\n self.num_heads = num_heads\n\n self.pe_embedder = EmbedND(dim=pe_dim, theta=self.theta, axes_dim=self.axes_dim)\n self.img_in = nn.Linear(in_channels, self.hidden_size, bias=True)\n\n self.txt_in = nn.Linear(context_in_dim, self.hidden_size)\n self.time_in = MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size)\n\n self.vector_in = MLPEmbedder(vec_in_dim, self.hidden_size)\n self.guidance_in = (MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size) if guidance_embed else nn.Identity())\n self.pos_embed_input = nn.Linear(in_channels, self.hidden_size, bias=True)\n self.gradient_checkpointing = False\n self.double_blocks = nn.ModuleList(\n [\n DoubleStreamBlock(\n self.hidden_size,\n self.num_heads,\n mlp_ratio=4.0,\n qkv_bias=True,\n )\n for _ in range(num_transformer)\n ]\n )\n\n self.single_blocks = nn.ModuleList(\n [\n SingleStreamBlock(self.hidden_size, self.num_heads, mlp_ratio=4.0)\n for _ in range(num_single_transformer)\n ]\n )\n\n # ControlNet blocks\n self.controlnet_blocks = nn.ModuleList([])\n for _ in range(num_transformer):\n controlnet_block = EnhanceControlnet(self.hidden_size)\n controlnet_block = zero_module(controlnet_block)\n self.controlnet_blocks.append(controlnet_block)\n\n # single controlnet blocks\n self.single_controlnet_blocks = nn.ModuleList([])\n for _ in range(num_single_transformer):\n controlnet_block = EnhanceControlnet(self.hidden_size)\n controlnet_block = zero_module(controlnet_block)\n self.single_controlnet_blocks.append(controlnet_block)\n\n # Input processing\n self.input_cond_block = CondDownsamplBlock()\n\n def _set_gradient_checkpointing(self, module, value=False):\n if hasattr(module, \"gradient_checkpointing\"):\n module.gradient_checkpointing = value\n\n\n @property\n def attn_processors(self):\n # set recursively\n processors = {}\n\n def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors):\n if hasattr(module, \"set_processor\"):\n processors[f\"{name}.processor\"] = module.processor\n\n for sub_name, child in module.named_children():\n fn_recursive_add_processors(f\"{name}.{sub_name}\", child, processors)\n\n return processors\n\n for name, module in self.named_children():\n fn_recursive_add_processors(name, module, processors)\n\n return processors\n\n def set_attn_processor(self, processor):\n r\"\"\"\n Sets the attention processor to use to compute attention.\n\n Parameters:\n processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`):\n The instantiated processor class or a dictionary of processor classes that will be set as the processor\n for **all** `Attention` layers.\n\n If `processor` is a dict, the key needs to define the path to the corresponding cross attention\n processor. This is strongly recommended when setting trainable attention processors.\n\n \"\"\"\n count = len(self.attn_processors.keys())\n\n if isinstance(processor, dict) and len(processor) != count:\n raise ValueError(\n f\"A dict of processors was passed, but the number of processors {len(processor)} does not match the\"\n f\" number of attention layers: {count}. Please make sure to pass {count} processor classes.\"\n )\n\n def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor):\n if hasattr(module, \"set_processor\"):\n if not isinstance(processor, dict):\n module.set_processor(processor)\n else:\n module.set_processor(processor.pop(f\"{name}.processor\"))\n\n for sub_name, child in module.named_children():\n fn_recursive_attn_processor(f\"{name}.{sub_name}\", child, processor)\n\n for name, module in self.named_children():\n fn_recursive_attn_processor(name, module, processor)\n\n def forward(\n self,\n img: Tensor,\n img_ids: Tensor,\n controlnet_cond: Tensor,\n txt: Tensor,\n txt_ids: Tensor,\n timesteps: Tensor,\n y: Tensor,\n guidance: Tensor | None = None,\n ) -> Tuple[Tensor, Tensor]:\n if img.ndim != 3 or txt.ndim != 3:\n raise ValueError(\"Input img and txt tensors must have 3 dimensions.\")\n\n # running on sequences img\n img = self.img_in(img)\n controlnet_cond = self.input_cond_block(controlnet_cond)\n controlnet_cond = rearrange(controlnet_cond, \"b c (h ph) (w pw) -> b (h w) (c ph pw)\", ph=2, pw=2)\n controlnet_cond = self.pos_embed_input(controlnet_cond)\n\n img = img + controlnet_cond\n\n vec = self.time_in(timestep_embedding(timesteps, 256))\n if self.guidance_embed:\n if guidance is None:\n raise ValueError(\"Didn't get guidance strength for guidance distilled model.\")\n vec = vec + self.guidance_in(timestep_embedding(guidance, 256))\n vec = vec + self.vector_in(y)\n txt = self.txt_in(txt)\n\n ids = torch.cat((txt_ids, img_ids), dim=1)\n pe = self.pe_embedder(ids)\n\n block_res_samples = ()\n for block in self.double_blocks:\n if self.training and self.gradient_checkpointing:\n def create_custom_forward(module, return_dict=None):\n def custom_forward(*inputs):\n if return_dict is not None:\n return module(*inputs, return_dict=return_dict)\n else:\n return module(*inputs)\n\n return custom_forward\n\n ckpt_kwargs: Dict[str, Any] = {\"use_reentrant\": False} if is_torch_version(\">=\", \"1.11.0\") else {}\n encoder_hidden_states, hidden_states = torch.utils.checkpoint.checkpoint(\n create_custom_forward(block),\n img,\n txt,\n vec,\n pe,\n )\n else:\n img, txt = block(img=img, txt=txt, vec=vec, pe=pe)\n block_res_samples = block_res_samples + (img,)\n\n img = torch.cat((txt, torch.zeros_like(img)), 1)\n single_block_res_samples = ()\n for index, block in enumerate(self.single_blocks):\n if self.training and self.gradient_checkpointing:\n def create_custom_forward(module, return_dict=None):\n def custom_forward(*inputs):\n if return_dict is not None:\n return module(*inputs, return_dict=return_dict)\n else:\n return module(*inputs)\n\n return custom_forward\n\n ckpt_kwargs: Dict[str, Any] = {\"use_reentrant\": False} if is_torch_version(\">=\", \"1.11.0\") else {}\n encoder_hidden_states, hidden_states = torch.utils.checkpoint.checkpoint(\n create_custom_forward(block),\n img,\n vec,\n pe,\n )\n else:\n img = block(img, vec=vec, pe=pe)\n single_block_res_samples = single_block_res_samples+(img,)\n\n controlnet_block_res_samples = ()\n for block_res_sample, controlnet_block in zip(block_res_samples, self.controlnet_blocks):\n block_res_sample = controlnet_block(block_res_sample)\n controlnet_block_res_samples = controlnet_block_res_samples + (block_res_sample,)\n\n single_controlnet_block_res_samples = ()\n for single_block_res_sample, single_controlnet_block in zip(single_block_res_samples, self.single_controlnet_blocks):\n single_block_res_sample = single_controlnet_block(single_block_res_sample)\n single_controlnet_block_res_samples = single_controlnet_block_res_samples + (single_block_res_sample,)\n\n return controlnet_block_res_samples,single_controlnet_block_res_samples"
},
{
"path": "modules/utils.py",
"content": "import math\nfrom dataclasses import dataclass\nfrom typing import Callable,Dict,Any\nimport torch\nfrom einops import rearrange\nfrom torch import Tensor, nn\nfrom tqdm import tqdm\n\n\ndef attention(q: Tensor, k: Tensor, v: Tensor, pe: Tensor) -> Tensor:\n q, k = apply_rope(q, k, pe)\n\n x = torch.nn.functional.scaled_dot_product_attention(q, k, v)\n x = rearrange(x, \"B H L D -> B L (H D)\")\n\n return x\n\n\ndef rope(pos: Tensor, dim: int, theta: int) -> Tensor:\n assert dim % 2 == 0\n scale = torch.arange(0, dim, 2, dtype=torch.float64, device=pos.device) / dim\n omega = 1.0 / (theta**scale)\n out = torch.einsum(\"...n,d->...nd\", pos, omega)\n out = torch.stack([torch.cos(out), -torch.sin(out), torch.sin(out), torch.cos(out)], dim=-1)\n out = rearrange(out, \"b n d (i j) -> b n d i j\", i=2, j=2)\n return out.float()\n\n\ndef apply_rope(xq: Tensor, xk: Tensor, freqs_cis: Tensor) -> tuple[Tensor, Tensor]:\n xq_ = xq.float().reshape(*xq.shape[:-1], -1, 1, 2)\n xk_ = xk.float().reshape(*xk.shape[:-1], -1, 1, 2)\n xq_out = freqs_cis[..., 0] * xq_[..., 0] + freqs_cis[..., 1] * xq_[..., 1]\n xk_out = freqs_cis[..., 0] * xk_[..., 0] + freqs_cis[..., 1] * xk_[..., 1]\n return xq_out.reshape(*xq.shape).type_as(xq), xk_out.reshape(*xk.shape).type_as(xk)\n\n\n\nclass EmbedND(nn.Module):\n def __init__(self, dim: int, theta: int, axes_dim: list[int]):\n super().__init__()\n self.dim = dim\n self.theta = theta\n self.axes_dim = axes_dim\n\n def forward(self, ids: Tensor) -> Tensor:\n n_axes = ids.shape[-1]\n emb = torch.cat(\n [rope(ids[..., i], self.axes_dim[i], self.theta) for i in range(n_axes)],\n dim=-3,\n )\n\n return emb.unsqueeze(1)\n\n\ndef timestep_embedding(t: Tensor, dim, max_period=10000, time_factor: float = 1000.0):\n \"\"\"\n Create sinusoidal timestep embeddings.\n :param t: a 1-D Tensor of N indices, one per batch element.\n These may be fractional.\n :param dim: the dimension of the output.\n :param max_period: controls the minimum frequency of the embeddings.\n :return: an (N, D) Tensor of positional embeddings.\n \"\"\"\n t = time_factor * t\n half = dim // 2\n freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to(\n t.device\n )\n\n args = t[:, None].float() * freqs[None]\n embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)\n if dim % 2:\n embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)\n if torch.is_floating_point(t):\n embedding = embedding.to(t)\n return embedding\n\n\nclass MLPEmbedder(nn.Module):\n def __init__(self, in_dim: int, hidden_dim: int):\n super().__init__()\n self.in_layer = nn.Linear(in_dim, hidden_dim, bias=True)\n self.silu = nn.SiLU()\n self.out_layer = nn.Linear(hidden_dim, hidden_dim, bias=True)\n\n def forward(self, x: Tensor) -> Tensor:\n return self.out_layer(self.silu(self.in_layer(x)))\n\n\nclass RMSNorm(torch.nn.Module):\n def __init__(self, dim: int):\n super().__init__()\n self.scale = nn.Parameter(torch.ones(dim))\n\n def forward(self, x: Tensor):\n x_dtype = x.dtype\n x = x.float()\n rrms = torch.rsqrt(torch.mean(x**2, dim=-1, keepdim=True) + 1e-6)\n return (x * rrms).to(dtype=x_dtype) * self.scale\n\n\nclass QKNorm(torch.nn.Module):\n def __init__(self, dim: int):\n super().__init__()\n self.query_norm = RMSNorm(dim)\n self.key_norm = RMSNorm(dim)\n\n def forward(self, q: Tensor, k: Tensor, v: Tensor) -> tuple[Tensor, Tensor]:\n q = self.query_norm(q)\n k = self.key_norm(k)\n return q.to(v), k.to(v)\n\n\nclass SelfAttention(nn.Module):\n def __init__(self, dim: int, num_heads: int = 8, qkv_bias: bool = False):\n super().__init__()\n self.num_heads = num_heads\n head_dim = dim // num_heads\n\n self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)\n self.norm = QKNorm(head_dim)\n self.proj = nn.Linear(dim, dim)\n\n def forward(self, x: Tensor, pe: Tensor) -> Tensor:\n qkv = self.qkv(x)\n q, k, v = rearrange(qkv, \"B L (K H D) -> K B H L D\", K=3, H=self.num_heads)\n q, k = self.norm(q, k, v)\n x = attention(q, k, v, pe=pe)\n x = self.proj(x)\n return x\n\n\n@dataclass\nclass ModulationOut:\n shift: Tensor\n scale: Tensor\n gate: Tensor\n\n\nclass Modulation(nn.Module):\n def __init__(self, dim: int, double: bool):\n super().__init__()\n self.is_double = double\n self.multiplier = 6 if double else 3\n self.lin = nn.Linear(dim, self.multiplier * dim, bias=True)\n\n def forward(self, vec: Tensor) -> tuple[ModulationOut, ModulationOut | None]:\n out = self.lin(nn.functional.silu(vec))[:, None, :].chunk(self.multiplier, dim=-1)\n\n return (\n ModulationOut(*out[:3]),\n ModulationOut(*out[3:]) if self.is_double else None,\n )\n\n\nclass DoubleStreamBlock(nn.Module):\n def __init__(self, hidden_size: int, num_heads: int, mlp_ratio: float, qkv_bias: bool = False):\n super().__init__()\n\n mlp_hidden_dim = int(hidden_size * mlp_ratio)\n self.num_heads = num_heads\n self.hidden_size = hidden_size\n self.img_mod = Modulation(hidden_size, double=True)\n self.img_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)\n self.img_attn = SelfAttention(dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias)\n\n self.img_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)\n self.img_mlp = nn.Sequential(\n nn.Linear(hidden_size, mlp_hidden_dim, bias=True),\n nn.GELU(approximate=\"tanh\"),\n nn.Linear(mlp_hidden_dim, hidden_size, bias=True),\n )\n\n self.txt_mod = Modulation(hidden_size, double=True)\n self.txt_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)\n self.txt_attn = SelfAttention(dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias)\n\n self.txt_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)\n self.txt_mlp = nn.Sequential(\n nn.Linear(hidden_size, mlp_hidden_dim, bias=True),\n nn.GELU(approximate=\"tanh\"),\n nn.Linear(mlp_hidden_dim, hidden_size, bias=True),\n )\n\n def forward(self, img: Tensor, txt: Tensor, vec: Tensor, pe: Tensor) -> tuple[Tensor, Tensor]:\n img_mod1, img_mod2 = self.img_mod(vec)\n txt_mod1, txt_mod2 = self.txt_mod(vec)\n\n # prepare image for attention\n img_modulated = self.img_norm1(img)\n img_modulated = (1 + img_mod1.scale) * img_modulated + img_mod1.shift\n img_qkv = self.img_attn.qkv(img_modulated)\n img_q, img_k, img_v = rearrange(img_qkv, \"B L (K H D) -> K B H L D\", K=3, H=self.num_heads)\n img_q, img_k = self.img_attn.norm(img_q, img_k, img_v)\n\n # prepare txt for attention\n txt_modulated = self.txt_norm1(txt)\n txt_modulated = (1 + txt_mod1.scale) * txt_modulated + txt_mod1.shift\n txt_qkv = self.txt_attn.qkv(txt_modulated)\n txt_q, txt_k, txt_v = rearrange(txt_qkv, \"B L (K H D) -> K B H L D\", K=3, H=self.num_heads)\n txt_q, txt_k = self.txt_attn.norm(txt_q, txt_k, txt_v)\n\n # run actual attention\n q = torch.cat((txt_q, img_q), dim=2)\n k = torch.cat((txt_k, img_k), dim=2)\n v = torch.cat((txt_v, img_v), dim=2)\n\n attn = attention(q, k, v, pe=pe)\n txt_attn, img_attn = attn[:, : txt.shape[1]], attn[:, txt.shape[1] :]\n\n # calculate the img bloks\n img = img + img_mod1.gate * self.img_attn.proj(img_attn)\n img = img + img_mod2.gate * self.img_mlp((1 + img_mod2.scale) * self.img_norm2(img) + img_mod2.shift)\n\n # calculate the txt bloks\n txt = txt + txt_mod1.gate * self.txt_attn.proj(txt_attn)\n txt = txt + txt_mod2.gate * self.txt_mlp((1 + txt_mod2.scale) * self.txt_norm2(txt) + txt_mod2.shift)\n return img, txt\n\n\nclass SingleStreamBlock(nn.Module):\n \"\"\"\n A DiT block with parallel linear layers as described in\n https://arxiv.org/abs/2302.05442 and adapted modulation interface.\n \"\"\"\n\n def __init__(\n self,\n hidden_size: int,\n num_heads: int,\n mlp_ratio: float = 4.0,\n qk_scale: float | None = None,\n ):\n super().__init__()\n self.hidden_dim = hidden_size\n self.num_heads = num_heads\n head_dim = hidden_size // num_heads\n self.scale = qk_scale or head_dim**-0.5\n\n self.mlp_hidden_dim = int(hidden_size * mlp_ratio)\n # qkv and mlp_in\n self.linear1 = nn.Linear(hidden_size, hidden_size * 3 + self.mlp_hidden_dim)\n # proj and mlp_out\n self.linear2 = nn.Linear(hidden_size + self.mlp_hidden_dim, hidden_size)\n\n self.norm = QKNorm(head_dim)\n\n self.hidden_size = hidden_size\n self.pre_norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)\n\n self.mlp_act = nn.GELU(approximate=\"tanh\")\n self.modulation = Modulation(hidden_size, double=False)\n\n def forward(self, x: Tensor, vec: Tensor, pe: Tensor) -> Tensor:\n mod, _ = self.modulation(vec)\n x_mod = (1 + mod.scale) * self.pre_norm(x) + mod.shift\n qkv, mlp = torch.split(self.linear1(x_mod), [3 * self.hidden_size, self.mlp_hidden_dim], dim=-1)\n\n q, k, v = rearrange(qkv, \"B L (K H D) -> K B H L D\", K=3, H=self.num_heads)\n q, k = self.norm(q, k, v)\n\n # compute attention\n attn = attention(q, k, v, pe=pe)\n # compute activation in mlp stream, cat again and run second linear layer\n output = self.linear2(torch.cat((attn, self.mlp_act(mlp)), 2))\n return x + mod.gate * output\n\n\nclass LastLayer(nn.Module):\n def __init__(self, hidden_size: int, patch_size: int, out_channels: int):\n super().__init__()\n self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)\n self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True)\n self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size, bias=True))\n\n def forward(self, x: Tensor, vec: Tensor) -> Tensor:\n shift, scale = self.adaLN_modulation(vec).chunk(2, dim=1)\n x = (1 + scale[:, None, :]) * self.norm_final(x) + shift[:, None, :]\n x = self.linear(x)\n return x\n\ndef get_noise(\n num_samples: int,\n height: int,\n width: int,\n device: torch.device,\n dtype: torch.dtype,\n seed: int,\n):\n return torch.randn(\n num_samples,\n 16,\n # allow for packing\n 2 * math.ceil(height / 16),\n 2 * math.ceil(width / 16),\n device=device,\n dtype=dtype,\n generator=torch.Generator(device=device).manual_seed(seed),\n )\n\ndef time_shift(mu: float, sigma: float, t: Tensor):\n return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma)\n\n\ndef get_lin_function(\n x1: float = 256, y1: float = 0.5, x2: float = 4096, y2: float = 1.15\n) -> Callable[[float], float]:\n m = (y2 - y1) / (x2 - x1)\n b = y1 - m * x1\n return lambda x: m * x + b\n\ndef get_schedule(\n num_steps: int,\n image_seq_len: int,\n base_shift: float = 0.5,\n max_shift: float = 1.15,\n shift: bool = True,\n) -> list[float]:\n # extra step for zero\n timesteps = torch.linspace(1, 0, num_steps + 1)\n\n # shifting the schedule to favor high timesteps for higher signal images\n if shift:\n # eastimate mu based on linear estimation between two points\n mu = get_lin_function(y1=base_shift, y2=max_shift)(image_seq_len)\n timesteps = time_shift(mu, 1.0, timesteps)\n\n return timesteps.tolist()\n\ndef unpack(x: Tensor, height: int, width: int) -> Tensor:\n return rearrange(\n x,\n \"b (h w) (c ph pw) -> b c (h ph) (w pw)\",\n h=math.ceil(height / 16),\n w=math.ceil(width / 16),\n ph=2,\n pw=2,\n )\n\ndef forward_mistoCN(\n model,\n img: Tensor,\n img_ids: Tensor,\n txt: Tensor,\n txt_ids: Tensor,\n timesteps: Tensor,\n y: Tensor,\n block_controlnet_hidden_states=None,\n single_controlnet_hidden_states=None,\n guidance: Tensor | None = None,\n):\n if img.ndim != 3 or txt.ndim != 3:\n raise ValueError(\"Input img and txt tensors must have 3 dimensions.\")\n\n # running on sequences img\n img = model.img_in(img)\n vec = model.time_in(timestep_embedding(timesteps, 256))\n vec = vec + model.guidance_in(timestep_embedding(guidance, 256))\n vec = vec + model.vector_in(y)\n txt = model.txt_in(txt)\n\n ids = torch.cat((txt_ids, img_ids), dim=1)\n pe = model.pe_embedder(ids)\n\n for index_block, block in enumerate(model.double_blocks):\n img, txt = block(img=img, txt=txt, vec=vec, pe=pe)\n # controlnet residual\n if block_controlnet_hidden_states is not None:\n if len(block_controlnet_hidden_states) == 1:\n img = img + block_controlnet_hidden_states[0]\n else:\n img = img + block_controlnet_hidden_states[index_block % 2]\n\n img = torch.cat((txt, img), 1)\n for index_block, block in enumerate(model.single_blocks):\n img = block(img, vec=vec, pe=pe)\n # controlnet residual\n if single_controlnet_hidden_states is not None:\n if len(single_controlnet_hidden_states) == 1:\n img = img + single_controlnet_hidden_states[0]\n else:\n img = img + single_controlnet_hidden_states[index_block % 2]\n\n img = img[:, txt.shape[1]:, ...]\n img = model.final_layer(img, vec) # (N, T, patch_size ** 2 * out_channels)\n return img\n\ndef denoise_controlnet(\n pbar,\n model,\n controlnet,\n # model input\n img: Tensor,\n img_ids: Tensor,\n txt: Tensor,\n txt_ids: Tensor,\n vec: Tensor,\n controlnet_cond,\n neg_txt,\n neg_txt_ids,\n neg_vec,\n # sampling parameters\n timesteps: list[float],\n guidance: float = 4.0,\n controlnet_strength=1.0,\n):\n controlnet.to(img.device, dtype=img.dtype)\n img_ids = img_ids.to(img.device, dtype=img.dtype)\n controlnet_cond = controlnet_cond.to(img.device, dtype=img.dtype)\n txt = txt.to(img.device, dtype=img.dtype)\n txt_ids = txt_ids.to(img.device, dtype=img.dtype)\n vec = vec.to(img.device, dtype=img.dtype)\n guidance_vec = torch.full((img.shape[0],), guidance, device=img.device, dtype=img.dtype)\n guidance_vec = guidance_vec.to(img.device, dtype=img.dtype)\n\n\n for t_curr, t_prev in tqdm(zip(timesteps[:-1], timesteps[1:]),desc=\"Sampling\", total = len(timesteps)-1):\n t_vec = torch.full((img.shape[0],), t_curr, dtype=img.dtype, device=img.device)\n block_res_samples, single_block_res_samples = controlnet(\n img=img,\n img_ids=img_ids,\n controlnet_cond=controlnet_cond,\n txt=txt,\n txt_ids=txt_ids,\n y=vec,\n timesteps=t_vec,\n guidance=guidance_vec,\n )\n pred = forward_mistoCN(\n model=model,\n img=img,\n img_ids=img_ids,\n txt=txt,\n txt_ids=txt_ids,\n y=vec,\n timesteps=t_vec,\n guidance=guidance_vec,\n block_controlnet_hidden_states=[i * controlnet_strength for i in block_res_samples],\n single_controlnet_hidden_states = [i * controlnet_strength for i in single_block_res_samples]\n )\n # negative\n neg_block_res_samples, neg_single_block_res_samples = controlnet(\n img=img,\n img_ids=img_ids,\n controlnet_cond=controlnet_cond,\n txt=neg_txt,\n txt_ids=neg_txt_ids,\n y=neg_vec,\n timesteps=t_vec,\n guidance=guidance_vec,\n )\n neg_pred = forward_mistoCN(\n model=model,\n img=img,\n img_ids=img_ids,\n txt=neg_txt,\n txt_ids=neg_txt_ids,\n y=neg_vec,\n timesteps=t_vec,\n guidance=guidance_vec,\n block_controlnet_hidden_states=[i * controlnet_strength for i in neg_block_res_samples],\n single_controlnet_hidden_states=[i * controlnet_strength for i in neg_single_block_res_samples]\n )\n f=0.85\n pred = neg_pred + f * (pred - neg_pred)\n img = img + (t_prev - t_curr) * pred\n pbar.update(1)\n\n return img"
},
{
"path": "nodes.py",
"content": "import torch\nimport os\nimport comfy.model_management\nfrom comfy.utils import ProgressBar\nimport folder_paths\nimport numpy as np\nfrom safetensors.torch import load_file\nfrom einops import rearrange,repeat\nfrom .modules.misto_controlnet import MistoControlNetFluxDev\nfrom .modules.utils import get_schedule,get_noise,denoise_controlnet, unpack\nimport torch.nn.functional as F\n\ndir_TheMistoModel = os.path.join(folder_paths.models_dir, \"TheMisto_model\")\nos.makedirs(dir_TheMistoModel, exist_ok=True)\nfolder_paths.folder_names_and_paths[\"TheMisto_model\"] = ([dir_TheMistoModel], folder_paths.supported_pt_extensions)\n\nclass LATENT_PROCESSOR_COMFY:\n def __init__(self):\n self.scale_factor = 0.3611\n self.shift_factor = 0.1159\n self.latent_rgb_factors =[\n [-0.0404, 0.0159, 0.0609],\n [ 0.0043, 0.0298, 0.0850],\n [ 0.0328, -0.0749, -0.0503],\n [-0.0245, 0.0085, 0.0549],\n [ 0.0966, 0.0894, 0.0530],\n [ 0.0035, 0.0399, 0.0123],\n [ 0.0583, 0.1184, 0.1262],\n [-0.0191, -0.0206, -0.0306],\n [-0.0324, 0.0055, 0.1001],\n [ 0.0955, 0.0659, -0.0545],\n [-0.0504, 0.0231, -0.0013],\n [ 0.0500, -0.0008, -0.0088],\n [ 0.0982, 0.0941, 0.0976],\n [-0.1233, -0.0280, -0.0897],\n [-0.0005, -0.0530, -0.0020],\n [-0.1273, -0.0932, -0.0680]\n ]\n def __call__(self, x):\n return (x / self.scale_factor) + self.shift_factor\n def go_back(self, x):\n return (x - self.shift_factor) * self.scale_factor\n\nMAX_RESOLUTION=16384\n\n\ndef prepare_sampling(t5_emb, clip_emb, img,batch_size):\n bs, c, h, w = img.shape\n bs = batch_size\n\n img = rearrange(img, \"b c (h ph) (w pw) -> b (h w) (c ph pw)\", ph=2, pw=2)\n if img.shape[0] == 1 and bs > 1:\n img = repeat(img, \"1 ... -> bs ...\", bs=bs)\n\n img_ids = torch.zeros(h // 2, w // 2, 3)\n img_ids[..., 1] = img_ids[..., 1] + torch.arange(h // 2)[:, None]\n img_ids[..., 2] = img_ids[..., 2] + torch.arange(w // 2)[None, :]\n img_ids = repeat(img_ids, \"h w c -> b (h w) c\", b=bs)\n\n if t5_emb.shape[0] == 1 and bs > 1:\n t5_emb = repeat(t5_emb, \"1 ... -> bs ...\", bs=bs)\n t5_emb_ids = torch.zeros(bs, t5_emb.shape[1], 3)\n\n if clip_emb.shape[0] == 1 and bs > 1:\n clip_emb = repeat(clip_emb, \"1 ... -> bs ...\", bs=bs)\n\n return {\n \"img\":img,\n \"img_ids\":img_ids.to(img.device, dtype=img.dtype),\n \"txt\":t5_emb.to(img.device, dtype=img.dtype),\n \"txt_ids\":t5_emb_ids.to(img.device, dtype=img.dtype),\n \"vec\":clip_emb.to(img.device, dtype=img.dtype)\n }\n\n\ndef load_misto_transoformer_cn(device):\n with torch.device(device):\n controlnet = MistoControlNetFluxDev(\n in_channels=64,\n vec_in_dim=768,\n context_in_dim=4096,\n hidden_size=3072,\n num_heads=24,\n num_transformer=3,\n num_single_transformer=2,\n guidance_embed=True,\n )\n return controlnet\n\n\ndef img_preprocessor(image, res):\n _, _, h, w = image.shape\n scale = res / min(h, w)\n new_h, new_w = int(h * scale), int(w * scale)\n resized = F.interpolate(image, size=(new_h, new_w), mode='bilinear', align_corners=False)\n crop_h = int((new_h // 16) * 16)\n crop_w = int((new_w // 16) * 16)\n start_h = (new_h - crop_h) // 2\n start_w = (new_w - crop_w) // 2\n cropped = resized[:, :, start_h:start_h + crop_h, start_w:start_w + crop_w]\n return cropped\n\n\n\nclass LoadMistoFluxControlNet:\n @classmethod\n def INPUT_TYPES(s):\n return {\"required\": {\n \"model_name\": (folder_paths.get_filename_list(\"TheMisto_model\"),)\n }}\n\n RETURN_TYPES = (\"MistoFluxControlNet\",)\n RETURN_NAMES = (\"ControlNet\",)\n FUNCTION = \"load_model\"\n CATEGORY = \"TheMistoAINodes\"\n\n def load_model(self,model_name):\n device=comfy.model_management.get_torch_device()\n misto_cn = load_misto_transoformer_cn(device=device)\n ckpt_path = os.path.join(dir_TheMistoModel, model_name)\n if '.bin' in model_name:\n state_dict = torch.load(ckpt_path, map_location='cpu')\n else:\n state_dict = load_file(ckpt_path)\n miss_, error_ = misto_cn.load_state_dict(state_dict,strict=False)\n misto_cn.eval()\n print(miss_, error_)\n return (misto_cn,)\n\nclass ApplyMistoFluxControlNet:\n @classmethod\n def INPUT_TYPES(s):\n return {\"required\": {\"controlnet\": (\"MistoFluxControlNet\",),\n \"image\": (\"IMAGE\",),\n \"resolution\": (\"INT\", {\"default\":960, \"min\": 512, \"max\": 4096}),\n \"strength\": (\"FLOAT\", {\"default\": 0.85, \"min\": 0.0, \"max\": 2.0, \"step\": 0.01})\n }}\n\n RETURN_TYPES = (\"ControlNetCondition\",\"IMAGE\")\n RETURN_NAMES = (\"controlnet_condition\",\"cond_image\")\n FUNCTION = \"embedding\"\n CATEGORY = \"TheMistoAINodes\"\n\n def embedding(self, controlnet, image, resolution, strength):\n cond_img = torch.from_numpy((np.array(image) * 2) - 1)\n cond_img = cond_img.permute(0, 3, 1, 2)\n res_img = img_preprocessor(image=cond_img, res=resolution)\n out_img = res_img.permute(0, 2, 3, 1)\n out_img = (out_img + 1) / 2\n cond_out = {\n \"img\": res_img,\n \"controlnet_strength\": strength,\n \"model\": controlnet,\n }\n return (cond_out,out_img)\n\n\nclass KSamplerTheMisto:\n @classmethod\n def INPUT_TYPES(s):\n return {\n \"required\": {\n \"model\": (\"MODEL\",),\n \"ae\":(\"VAE\",),\n \"positive\": (\"CONDITIONING\",),\n \"negative\": (\"CONDITIONING\",),\n \"controlnet_condition\": (\"ControlNetCondition\", {\"default\": None}),\n \"batch_size\": (\"INT\", {\"default\":1, \"min\": 1, \"max\": 100}),\n \"guidance\": (\"FLOAT\", {\"default\": 3.5, \"min\": 0.1, \"max\": 30}),\n \"seed\": (\"INT\", {\"default\": 0, \"min\": 0, \"max\": 0xffffffffffffffff}),\n \"steps\": (\"INT\", {\"default\": 20, \"min\": 1, \"max\": 100}),\n },\n }\n\n RETURN_TYPES = (\"IMAGE\",)\n RETURN_NAMES = (\"image\",)\n FUNCTION = \"sampling\"\n CATEGORY = \"TheMistoAINodes\"\n\n def sampling(self, model,ae, positive, negative,controlnet_condition,batch_size,guidance,seed, steps ):\n # device ,dtype and pbar\n device = comfy.model_management.get_torch_device()\n dtype_model = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16\n pbar = ProgressBar(steps+10)\n pbar.update(1)\n\n # model\n comfy.model_management.load_model_gpu(model)\n flux_model = model.model.diffusion_model\n pbar.update(3)\n\n # cn cond\n cn_model = controlnet_condition['model']\n cond_img = controlnet_condition['img'].to(torch.bfloat16).to(device)\n cn_strength = controlnet_condition['controlnet_strength']\n\n bc, c, h, w = cond_img.shape\n height = (h//16) * 16\n width = (w//16) * 16\n pbar.update(2)\n with torch.no_grad():\n # set scheduler\n timesteps = get_schedule( steps, (width // 8) * (height // 8) // 4, shift=True, )\n x = get_noise( 1, height, width, device=device, dtype=dtype_model, seed=seed)\n p_inp_cond = prepare_sampling(positive[0][0], positive[0][1]['pooled_output'], img=x, batch_size=batch_size)\n n_inp_cond = prepare_sampling(negative[0][0], negative[0][1]['pooled_output'], img=x, batch_size=batch_size)\n pbar.update(1)\n # denoise\n x = denoise_controlnet(\n pbar=pbar,\n model=flux_model, **p_inp_cond,\n controlnet=cn_model,\n timesteps=timesteps,\n guidance=guidance,\n controlnet_cond=cond_img,\n controlnet_strength = cn_strength,\n neg_txt=n_inp_cond['txt'],\n neg_txt_ids=n_inp_cond['txt_ids'],\n neg_vec=n_inp_cond['vec'],\n )\n\n x = unpack(x.float(), height, width)\n lat_processor = LATENT_PROCESSOR_COMFY()\n x = lat_processor(x)\n pbar.update(1)\n return (ae.decode(x),)\n\nNODE_CLASS_MAPPINGS = {\n \"LoadTheMistoFluxControlNet\": LoadMistoFluxControlNet,\n \"ApplyTheMistoFluxControlNet\": ApplyMistoFluxControlNet,\n \"KSamplerTheMisto\":KSamplerTheMisto,\n}\nNODE_DISPLAY_NAME_MAPPINGS = {\n \"LoadTheMistoFluxControlNet\": \"Load MistoCN-Flux.dev\",\n \"ApplyTheMistoFluxControlNet\": \"Apply MistoCN-Flux.dev\",\n \"KSamplerTheMisto\":\"KSampler for MistoCN-Flux.dev\",\n}\n"
},
{
"path": "workflows/example_workflow.json",
"content": "{\n \"last_node_id\": 40,\n \"last_link_id\": 44,\n \"nodes\": [\n {\n \"id\": 13,\n \"type\": \"VAELoader\",\n \"pos\": {\n \"0\": 60,\n \"1\": 327,\n \"2\": 0,\n \"3\": 0,\n \"4\": 0,\n \"5\": 0,\n \"6\": 0,\n \"7\": 0,\n \"8\": 0,\n \"9\": 0\n },\n \"size\": {\n \"0\": 315,\n \"1\": 58\n },\n \"flags\": {},\n \"order\": 0,\n \"mode\": 0,\n \"inputs\": [],\n \"outputs\": [\n {\n \"name\": \"VAE\",\n \"type\": \"VAE\",\n \"links\": [\n 19\n ],\n \"slot_index\": 0,\n \"shape\": 3\n }\n ],\n \"properties\": {\n \"Node name for S&R\": \"VAELoader\"\n },\n \"widgets_values\": [\n \"ae.safetensors\"\n ]\n },\n {\n \"id\": 10,\n \"type\": \"DualCLIPLoader\",\n \"pos\": {\n \"0\": 58,\n \"1\": 27,\n \"2\": 0,\n \"3\": 0,\n \"4\": 0,\n \"5\": 0,\n \"6\": 0,\n \"7\": 0,\n \"8\": 0,\n \"9\": 0\n },\n \"size\": {\n \"0\": 348.2605285644531,\n \"1\": 106\n },\n \"flags\": {},\n \"order\": 1,\n \"mode\": 0,\n \"inputs\": [],\n \"outputs\": [\n {\n \"name\": \"CLIP\",\n \"type\": \"CLIP\",\n \"links\": [\n 24,\n 26\n ],\n \"slot_index\": 0,\n \"shape\": 3\n }\n ],\n \"properties\": {\n \"Node name for S&R\": \"DualCLIPLoader\"\n },\n \"widgets_values\": [\n \"clip_l.safetensors\",\n \"t5xxl_fp8_e4m3fn.safetensors\",\n \"flux\"\n ]\n },\n {\n \"id\": 36,\n \"type\": \"AnyLineArtPreprocessor_aux\",\n \"pos\": {\n \"0\": 250,\n \"1\": 569,\n \"2\": 0,\n \"3\": 0,\n \"4\": 0,\n \"5\": 0,\n \"6\": 0,\n \"7\": 0,\n \"8\": 0,\n \"9\": 0\n },\n \"size\": {\n \"0\": 315,\n \"1\": 178\n },\n \"flags\": {},\n \"order\": 11,\n \"mode\": 0,\n \"inputs\": [\n {\n \"name\": \"image\",\n \"type\": \"IMAGE\",\n \"link\": 38\n }\n ],\n \"outputs\": [\n {\n \"name\": \"image\",\n \"type\": \"IMAGE\",\n \"links\": [\n 37\n ],\n \"shape\": 3,\n \"slot_index\": 0\n }\n ],\n \"properties\": {\n \"Node name for S&R\": \"AnyLineArtPreprocessor_aux\"\n },\n \"widgets_values\": [\n \"lineart_standard\",\n 960,\n 0,\n 1,\n 36,\n 1\n ]\n },\n {\n \"id\": 17,\n \"type\": \"LoadTheMistoFluxControlNet\",\n \"pos\": {\n \"0\": 54,\n \"1\": 445,\n \"2\": 0,\n \"3\": 0,\n \"4\": 0,\n \"5\": 0,\n \"6\": 0,\n \"7\": 0,\n \"8\": 0,\n \"9\": 0\n },\n \"size\": {\n \"0\": 392,\n \"1\": 60\n },\n \"flags\": {},\n \"order\": 2,\n \"mode\": 0,\n \"inputs\": [],\n \"outputs\": [\n {\n \"name\": \"ControlNet\",\n \"type\": \"MistoFluxControlNet\",\n \"links\": [\n 40\n ],\n \"slot_index\": 0,\n \"shape\": 3\n }\n ],\n \"properties\": {\n \"Node name for S&R\": \"LoadTheMistoFluxControlNet\"\n },\n \"widgets_values\": [\n \"mistoline_flux.dev_v1.safetensors\"\n ]\n },\n {\n \"id\": 37,\n \"type\": \"LoadImage\",\n \"pos\": {\n \"0\": -92,\n \"1\": 568,\n \"2\": 0,\n \"3\": 0,\n \"4\": 0,\n \"5\": 0,\n \"6\": 0,\n \"7\": 0,\n \"8\": 0,\n \"9\": 0\n },\n \"size\": [\n 315,\n 314\n ],\n \"flags\": {},\n \"order\": 3,\n \"mode\": 0,\n \"inputs\": [],\n \"outputs\": [\n {\n \"name\": \"IMAGE\",\n \"type\": \"IMAGE\",\n \"links\": [\n 38\n ],\n \"shape\": 3,\n \"slot_index\": 0\n },\n {\n \"name\": \"MASK\",\n \"type\": \"MASK\",\n \"links\": null,\n \"shape\": 3\n }\n ],\n \"properties\": {\n \"Node name for S&R\": \"LoadImage\"\n },\n \"widgets_values\": [\n \"image--73-.png\",\n \"image\"\n ]\n },\n {\n \"id\": 26,\n \"type\": \"CLIPTextEncodeFlux\",\n \"pos\": {\n \"0\": 819,\n \"1\": 29,\n \"2\": 0,\n \"3\": 0,\n \"4\": 0,\n \"5\": 0,\n \"6\": 0,\n \"7\": 0,\n \"8\": 0,\n \"9\": 0\n },\n \"size\": {\n \"0\": 400,\n \"1\": 200\n },\n \"flags\": {},\n \"order\": 9,\n \"mode\": 0,\n \"inputs\": [\n {\n \"name\": \"clip\",\n \"type\": \"CLIP\",\n \"link\": 24\n }\n ],\n \"outputs\": [\n {\n \"name\": \"CONDITIONING\",\n \"type\": \"CONDITIONING\",\n \"links\": [\n 25\n ],\n \"slot_index\": 0,\n \"shape\": 3\n }\n ],\n \"properties\": {\n \"Node name for S&R\": \"CLIPTextEncodeFlux\"\n },\n \"widgets_values\": [\n \"Model in modern attire with metallic accessories, in an old factory setting, Metallic sheen, Full-frame mirrorless, 35mm lens, f/2.8 aperture, ISO 500, off-camera flash\",\n \"Model in modern attire with metallic accessories, in an old factory setting, Metallic sheen, Full-frame mirrorless, 35mm lens, f/2.8 aperture, ISO 500, off-camera flash\",\n 3.5\n ]\n },\n {\n \"id\": 4,\n \"type\": \"CheckpointLoaderSimple\",\n \"pos\": {\n \"0\": 64,\n \"1\": 179,\n \"2\": 0,\n \"3\": 0,\n \"4\": 0,\n \"5\": 0,\n \"6\": 0,\n \"7\": 0,\n \"8\": 0,\n \"9\": 0\n },\n \"size\": {\n \"0\": 315,\n \"1\": 98\n },\n \"flags\": {},\n \"order\": 4,\n \"mode\": 0,\n \"inputs\": [],\n \"outputs\": [\n {\n \"name\": \"MODEL\",\n \"type\": \"MODEL\",\n \"links\": [\n 12\n ],\n \"slot_index\": 0\n },\n {\n \"name\": \"CLIP\",\n \"type\": \"CLIP\",\n \"links\": [],\n \"slot_index\": 1\n },\n {\n \"name\": \"VAE\",\n \"type\": \"VAE\",\n \"links\": [],\n \"slot_index\": 2\n }\n ],\n \"properties\": {\n \"Node name for S&R\": \"CheckpointLoaderSimple\"\n },\n \"widgets_values\": [\n \"flux1-dev-fp8.safetensors\"\n ]\n },\n {\n \"id\": 28,\n \"type\": \"CLIPTextEncodeFlux\",\n \"pos\": {\n \"0\": 823,\n \"1\": 273,\n \"2\": 0,\n \"3\": 0,\n \"4\": 0,\n \"5\": 0,\n \"6\": 0,\n \"7\": 0,\n \"8\": 0,\n \"9\": 0\n },\n \"size\": {\n \"0\": 399.6443786621094,\n \"1\": 156.8167266845703\n },\n \"flags\": {},\n \"order\": 10,\n \"mode\": 0,\n \"inputs\": [\n {\n \"name\": \"clip\",\n \"type\": \"CLIP\",\n \"link\": 26\n }\n ],\n \"outputs\": [\n {\n \"name\": \"CONDITIONING\",\n \"type\": \"CONDITIONING\",\n \"links\": [\n 27\n ],\n \"slot_index\": 0,\n \"shape\": 3\n }\n ],\n \"properties\": {\n \"Node name for S&R\": \"CLIPTextEncodeFlux\"\n },\n \"widgets_values\": [\n \"out of frame, lowres, text, error, cropped, worst quality, low quality, jpeg artifacts, ugly, duplicate, morbid, mutilated, out of frame, extra fingers, mutated hands, poorly drawn hands, poorly drawn face, mutation, deformed, blurry, bad anatomy, bad proportions, extra limbs, cloned face, disfigured, gross proportions, malformed limbs, missing arms, missing legs, extra arms, extra legs, fused fingers, too many fingers, long neck, username, watermark, signature.\",\n \"out of frame, lowres, text, error, cropped, worst quality, low quality, jpeg artifacts, ugly, duplicate, morbid, mutilated, out of frame, extra fingers, mutated hands, poorly drawn hands, poorly drawn face, mutation, deformed, blurry, bad anatomy, bad proportions, extra limbs, cloned face, disfigured, gross proportions, malformed limbs, missing arms, missing legs, extra arms, extra legs, fused fingers, too many fingers, long neck, username, watermark, signature.\",\n 3.5\n ]\n },\n {\n \"id\": 21,\n \"type\": \"Note\",\n \"pos\": {\n \"0\": 587,\n \"1\": 286,\n \"2\": 0,\n \"3\": 0,\n \"4\": 0,\n \"5\": 0,\n \"6\": 0,\n \"7\": 0,\n \"8\": 0,\n \"9\": 0\n },\n \"size\": {\n \"0\": 210,\n \"1\": 122.03421020507812\n },\n \"flags\": {},\n \"order\": 5,\n \"mode\": 0,\n \"inputs\": [],\n \"outputs\": [],\n \"properties\": {\n \"text\": \"\"\n },\n \"widgets_values\": [\n \"Make sure to craft your prompts well—precision is more important than length, otherwise the result could be very bad.\\n\\nPrompt非常重要，请用AI或者自己写一个好的prompt，不然效果会很糟糕。\"\n ],\n \"color\": \"#432\",\n \"bgcolor\": \"#653\"\n },\n {\n \"id\": 39,\n \"type\": \"Note\",\n \"pos\": {\n \"0\": 774,\n \"1\": 629,\n \"2\": 0,\n \"3\": 0,\n \"4\": 0,\n \"5\": 0,\n \"6\": 0,\n \"7\": 0,\n \"8\": 0,\n \"9\": 0\n },\n \"size\": [\n 210,\n 122.03421020507812\n ],\n \"flags\": {},\n \"order\": 6,\n \"mode\": 0,\n \"inputs\": [],\n \"outputs\": [],\n \"properties\": {\n \"text\": \"\"\n },\n \"widgets_values\": [\n \"Very sensitive to strength, please start with about 0.5 and gradually increase\\n\\n对strength非常敏感，请从0.5开始尝试逐步增加。\"\n ],\n \"color\": \"#432\",\n \"bgcolor\": \"#653\"\n },\n {\n \"id\": 38,\n \"type\": \"PreviewImage\",\n \"pos\": {\n \"0\": 1015,\n \"1\": 479,\n \"2\": 0,\n \"3\": 0,\n \"4\": 0,\n \"5\": 0,\n \"6\": 0,\n \"7\": 0,\n \"8\": 0,\n \"9\": 0\n },\n \"size\": [\n 332.628470953563,\n 319.48089726588853\n ],\n \"flags\": {},\n \"order\": 15,\n \"mode\": 0,\n \"inputs\": [\n {\n \"name\": \"images\",\n \"type\": \"IMAGE\",\n \"link\": 41\n }\n ],\n \"outputs\": [],\n \"properties\": {\n \"Node name for S&R\": \"PreviewImage\"\n }\n },\n {\n \"id\": 35,\n \"type\": \"ApplyTheMistoFluxControlNet\",\n \"pos\": {\n \"0\": 593,\n \"1\": 482,\n \"2\": 0,\n \"3\": 0,\n \"4\": 0,\n \"5\": 0,\n \"6\": 0,\n \"7\": 0,\n \"8\": 0,\n \"9\": 0\n },\n \"size\": {\n \"0\": 393,\n \"1\": 102\n },\n \"flags\": {},\n \"order\": 13,\n \"mode\": 0,\n \"inputs\": [\n {\n \"name\": \"controlnet\",\n \"type\": \"MistoFluxControlNet\",\n \"link\": 40\n },\n {\n \"name\": \"image\",\n \"type\": \"IMAGE\",\n \"link\": 37\n }\n ],\n \"outputs\": [\n {\n \"name\": \"controlnet_condition\",\n \"type\": \"ControlNetCondition\",\n \"links\": [\n 39\n ],\n \"shape\": 3,\n \"slot_index\": 0\n },\n {\n \"name\": \"cond_image\",\n \"type\": \"IMAGE\",\n \"links\": [\n 41\n ],\n \"shape\": 3,\n \"slot_index\": 1\n }\n ],\n \"properties\": {\n \"Node name for S&R\": \"ApplyTheMistoFluxControlNet\"\n },\n \"widgets_values\": [\n 960,\n 0.5\n ]\n },\n {\n \"id\": 15,\n \"type\": \"LoraLoaderModelOnly\",\n \"pos\": {\n \"0\": 498,\n \"1\": -118,\n \"2\": 0,\n \"3\": 0,\n \"4\": 0,\n \"5\": 0,\n \"6\": 0,\n \"7\": 0,\n \"8\": 0,\n \"9\": 0\n },\n \"size\": {\n \"0\": 380.00384521484375,\n \"1\": 82\n },\n \"flags\": {},\n \"order\": 12,\n \"mode\": 0,\n \"inputs\": [\n {\n \"name\": \"model\",\n \"type\": \"MODEL\",\n \"link\": 12\n }\n ],\n \"outputs\": [\n {\n \"name\": \"MODEL\",\n \"type\": \"MODEL\",\n \"links\": [\n 44\n ],\n \"slot_index\": 0,\n \"shape\": 3\n }\n ],\n \"properties\": {\n \"Node name for S&R\": \"LoraLoaderModelOnly\"\n },\n \"widgets_values\": [\n \"Hyper-FLUX.1-dev-16steps-lora.safetensors\",\n 0.13\n ]\n },\n {\n \"id\": 19,\n \"type\": \"KSamplerTheMisto\",\n \"pos\": {\n \"0\": 1251,\n \"1\": -121,\n \"2\": 0,\n \"3\": 0,\n \"4\": 0,\n \"5\": 0,\n \"6\": 0,\n \"7\": 0,\n \"8\": 0,\n \"9\": 0\n },\n \"size\": {\n \"0\": 330,\n \"1\": 234\n },\n \"flags\": {},\n \"order\": 14,\n \"mode\": 0,\n \"inputs\": [\n {\n \"name\": \"model\",\n \"type\": \"MODEL\",\n \"link\": 44\n },\n {\n \"name\": \"ae\",\n \"type\": \"VAE\",\n \"link\": 19\n },\n {\n \"name\": \"positive\",\n \"type\": \"CONDITIONING\",\n \"link\": 25\n },\n {\n \"name\": \"negative\",\n \"type\": \"CONDITIONING\",\n \"link\": 27\n },\n {\n \"name\": \"controlnet_condition\",\n \"type\": \"ControlNetCondition\",\n \"link\": 39\n }\n ],\n \"outputs\": [\n {\n \"name\": \"image\",\n \"type\": \"IMAGE\",\n \"links\": [\n 20\n ],\n \"slot_index\": 0,\n \"shape\": 3\n }\n ],\n \"properties\": {\n \"Node name for S&R\": \"KSamplerTheMisto\"\n },\n \"widgets_values\": [\n 1,\n 4,\n 1039574509141705,\n \"randomize\",\n 16\n ]\n },\n {\n \"id\": 22,\n \"type\": \"PreviewImage\",\n \"pos\": {\n \"0\": 1614,\n \"1\": -128,\n \"2\": 0,\n \"3\": 0,\n \"4\": 0,\n \"5\": 0,\n \"6\": 0,\n \"7\": 0,\n \"8\": 0,\n \"9\": 0\n },\n \"size\": [\n 461.88277135920316,\n 499.946617357832\n ],\n \"flags\": {},\n \"order\": 16,\n \"mode\": 0,\n \"inputs\": [\n {\n \"name\": \"images\",\n \"type\": \"IMAGE\",\n \"link\": 20\n }\n ],\n \"outputs\": [],\n \"properties\": {\n \"Node name for S&R\": \"PreviewImage\"\n }\n },\n {\n \"id\": 20,\n \"type\": \"Note\",\n \"pos\": {\n \"0\": 586,\n \"1\": 55,\n \"2\": 0,\n \"3\": 0,\n \"4\": 0,\n \"5\": 0,\n \"6\": 0,\n \"7\": 0,\n \"8\": 0,\n \"9\": 0\n },\n \"size\": {\n \"0\": 210,\n \"1\": 122.03421020507812\n },\n \"flags\": {},\n \"order\": 7,\n \"mode\": 0,\n \"inputs\": [],\n \"outputs\": [],\n \"properties\": {\n \"text\": \"\"\n },\n \"widgets_values\": [\n \"Make sure to craft your prompts well—precision is more important than length, otherwise the result could be very bad.\\n\\nPrompt非常重要，请用AI或者自己写一个好的prompt，不然效果会很糟糕。\"\n ],\n \"color\": \"#432\",\n \"bgcolor\": \"#653\"\n },\n {\n \"id\": 40,\n \"type\": \"Note\",\n \"pos\": {\n \"0\": 1258,\n \"1\": -232,\n \"2\": 0,\n \"3\": 0,\n \"4\": 0,\n \"5\": 0,\n \"6\": 0,\n \"7\": 0,\n \"8\": 0,\n \"9\": 0\n },\n \"size\": [\n 210,\n 67.63623725377664\n ],\n \"flags\": {},\n \"order\": 8,\n \"mode\": 0,\n \"inputs\": [],\n \"outputs\": [],\n \"properties\": {\n \"text\": \"\"\n },\n \"widgets_values\": [\n \"no autoCFG\\n\\n不要使用自动CFG\"\n ],\n \"color\": \"#432\",\n \"bgcolor\": \"#653\"\n }\n ],\n \"links\": [\n [\n 12,\n 4,\n 0,\n 15,\n 0,\n \"MODEL\"\n ],\n [\n 19,\n 13,\n 0,\n 19,\n 1,\n \"VAE\"\n ],\n [\n 20,\n 19,\n 0,\n 22,\n 0,\n \"IMAGE\"\n ],\n [\n 24,\n 10,\n 0,\n 26,\n 0,\n \"CLIP\"\n ],\n [\n 25,\n 26,\n 0,\n 19,\n 2,\n \"CONDITIONING\"\n ],\n [\n 26,\n 10,\n 0,\n 28,\n 0,\n \"CLIP\"\n ],\n [\n 27,\n 28,\n 0,\n 19,\n 3,\n \"CONDITIONING\"\n ],\n [\n 37,\n 36,\n 0,\n 35,\n 1,\n \"IMAGE\"\n ],\n [\n 38,\n 37,\n 0,\n 36,\n 0,\n \"IMAGE\"\n ],\n [\n 39,\n 35,\n 0,\n 19,\n 4,\n \"ControlNetCondition\"\n ],\n [\n 40,\n 17,\n 0,\n 35,\n 0,\n \"MistoFluxControlNet\"\n ],\n [\n 41,\n 35,\n 1,\n 38,\n 0,\n \"IMAGE\"\n ],\n [\n 44,\n 15,\n 0,\n 19,\n 0,\n \"MODEL\"\n ]\n ],\n \"groups\": [],\n \"config\": {},\n \"extra\": {\n \"ds\": {\n \"scale\": 0.9090909090909091,\n \"offset\": [\n 384.31753255840897,\n 385.7533316712783\n ]\n }\n },\n \"version\": 0.4\n}"
}
]