init
This commit is contained in:
smthemex
+1
-9
@@ -1,10 +1,2 @@
|
||||
script/
|
||||
*.symlink
|
||||
output/
|
||||
regress_result/
|
||||
google
|
||||
checkpoints
|
||||
huggingface_models
|
||||
openai
|
||||
__pycache__/
|
||||
.vscode/
|
||||
|
||||
|
||||
@@ -1,27 +1,21 @@
|
||||
# ComfyUI_KV_Edit
|
||||
[KV_Edit](https://github.com/Xilluill/KV-Edit): Training-Free Image Editing for Precise Background Preservation，you can use it in comfyUI
|
||||
|
||||
[KV_Edit](https://github.com/Xilluill/KV-Edit): Training-Free Image Editing for Precise
|
||||
Background Preservation, re-implemented **on ComfyUI's own FLUX DiT stack**.
|
||||
|
||||
# update
|
||||
* use single dit like kj or other /修改模型加载占用及去掉重复量化，加入加速lora支持
|
||||
* 改成comfyUI的内核封装 /changed to ComfyUI's own FLUX DiT stack
|
||||
|
||||
|
||||
# 1. Installation
|
||||
|
||||
In the ./ComfyUI/custom_node directory, run the following:
|
||||
In the ./ComfyUI/custom_nodes directory, run the following:
|
||||
```
|
||||
git clone https://github.com/smthemex/ComfyUI_KV_Edit
|
||||
```
|
||||
---
|
||||
|
||||
# 2. Requirements
|
||||
|
||||
```
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
# 3.Model
|
||||
* 3.1.1 download 'ae.safetensor' and 'flux1-dev.safetensors' from [here](https://huggingface.co/black-forest-labs/FLUX.1-dev/tree/main) 下载ae和flux单体模型模型,文件结构如下图
|
||||
* 3.1.2 download 'clip_l.safetensors' and 't5xxl_fp8_e4m3fn.safetensors' from [here](https://huggingface.co/comfyanonymous/flux_text_encoders/tree/main)下载T5和clip l单体模型模型,文件结构如下图
|
||||
# 2.Model
|
||||
* 2.1 download 'ae.safetensor' and 'flux1-dev.safetensors' from [here](https://huggingface.co/black-forest-labs/FLUX.1-dev/tree/main) 下载ae和flux单体模型模型,文件结构如下图
|
||||
* 2.2 download 'clip_l.safetensors' and 't5xxl_fp8_e4m3fn.safetensors' from [here](https://huggingface.co/comfyanonymous/flux_text_encoders/tree/main)下载T5和clip l单体模型模型,文件结构如下图
|
||||
|
||||
```
|
||||
-- ComfyUI/models/diffusion_models # or unet
|
||||
@@ -34,19 +28,12 @@ pip install -r requirements.txt
|
||||
├── ae.safetensor
|
||||
```
|
||||
|
||||
# Example
|
||||
# 3 Example
|
||||
* new
|
||||
|
||||
* use comfy T5 and clip and flux dev single checkpoints用comfy的text encoder和flux的fp16/fp8模型，推荐
|
||||
|
||||
* use comfy or flux dev single checkpoints 仅使用单体模型，速度更慢，更占内存，主要是量化两次模型
|
||||
|
||||
|
||||
|
||||
# 4 Citation
|
||||
|
||||
|
||||
# Citation
|
||||
|
||||
If you find our work helpful, please **star 🌟** this[ repo](https://github.com/Xilluill/KV-Edit) and **cite 📑** our paper. Thanks for your support!
|
||||
```
|
||||
@article{zhu2025kv,
|
||||
title={KV-Edit: Training-Free Image Editing for Precise Background Preservation},
|
||||
@@ -55,9 +42,3 @@ If you find our work helpful, please **star 🌟** this[ repo](https://github.co
|
||||
year={2025}
|
||||
}
|
||||
```
|
||||
|
||||
# 👍🏻 Acknowledgements
|
||||
Our code is modified based on [FLUX](https://github.com/black-forest-labs/flux) and [RF-Solver-Edit](https://github.com/wangjiangshan0725/RF-Solver-Edit). Special thanks to [Wenke Huang](https://wenkehuang.github.io/) for his early inspiration and helpful guidance to this project!
|
||||
|
||||
# 📧 Contact
|
||||
This repository is currently under active development and restructuring. The codebase is being optimized for better stability and reproducibility. While we strive to maintain code quality, you may encounter temporary issues during this transition period. For any questions or technical discussions, feel free to open an issue or contact us via email at xilluill070513@gmail.com.
|
||||
|
||||
+13
-1
@@ -1,4 +1,16 @@
|
||||
# -*- coding: UTF-8 -*-
|
||||
"""ComfyUI_KV_Edit - native-diT rewrite.
|
||||
|
||||
from .kv_edit_node import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
|
||||
The old wrapper (``kv_edit_node.py`` + ``flux/`` + ``gradio_*.py``) is still on disk for
|
||||
reference but is **not imported any more**: it built a private copy of the FLUX model
|
||||
(next to the one comfy already had), loaded it through ``accelerate`` and reached into
|
||||
``huggingface_hub``, so every diffusers / transformers / accelerate bump could break it
|
||||
and it doubled the VRAM.
|
||||
|
||||
The native path needs nothing beyond torch: the diffusion model, the CLIP, the VAE and
|
||||
the sampler all stay in their own nodes.
|
||||
"""
|
||||
|
||||
from .kvedit_nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
|
||||
|
||||
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
|
||||
|
||||
Binary file not shown.
|
Before Width: | Height: | Size: 222 KiB After Width: | Height: | Size: 580 KiB |
@@ -1,592 +0,0 @@
|
||||
{
|
||||
"id": "6e3edd3c-6fa2-4d01-ad55-33e37446634d",
|
||||
"revision": 0,
|
||||
"last_node_id": 34,
|
||||
"last_link_id": 70,
|
||||
"nodes": [
|
||||
{
|
||||
"id": 14,
|
||||
"type": "SaveImage",
|
||||
"pos": [
|
||||
21780.08203125,
|
||||
-884.1846923828125
|
||||
],
|
||||
"size": [
|
||||
305.730712890625,
|
||||
292.5241394042969
|
||||
],
|
||||
"flags": {},
|
||||
"order": 11,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "images",
|
||||
"type": "IMAGE",
|
||||
"link": 46
|
||||
}
|
||||
],
|
||||
"outputs": [],
|
||||
"properties": {},
|
||||
"widgets_values": [
|
||||
"ComfyUI"
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 31,
|
||||
"type": "MaskToImage",
|
||||
"pos": [
|
||||
20504.787109375,
|
||||
-440.2896423339844
|
||||
],
|
||||
"size": [
|
||||
140,
|
||||
26
|
||||
],
|
||||
"flags": {},
|
||||
"order": 6,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "mask",
|
||||
"type": "MASK",
|
||||
"link": 65
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "IMAGE",
|
||||
"type": "IMAGE",
|
||||
"links": [
|
||||
66
|
||||
]
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "MaskToImage"
|
||||
},
|
||||
"widgets_values": []
|
||||
},
|
||||
{
|
||||
"id": 32,
|
||||
"type": "ImageScale",
|
||||
"pos": [
|
||||
20721.037109375,
|
||||
-276.83343505859375
|
||||
],
|
||||
"size": [
|
||||
270,
|
||||
130
|
||||
],
|
||||
"flags": {},
|
||||
"order": 7,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "image",
|
||||
"type": "IMAGE",
|
||||
"link": 66
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "IMAGE",
|
||||
"type": "IMAGE",
|
||||
"links": [
|
||||
67
|
||||
]
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "ImageScale"
|
||||
},
|
||||
"widgets_values": [
|
||||
"nearest-exact",
|
||||
512,
|
||||
512,
|
||||
"disabled"
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 30,
|
||||
"type": "ImageScale",
|
||||
"pos": [
|
||||
20682.861328125,
|
||||
-511.3575744628906
|
||||
],
|
||||
"size": [
|
||||
270,
|
||||
130
|
||||
],
|
||||
"flags": {},
|
||||
"order": 5,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "image",
|
||||
"type": "IMAGE",
|
||||
"link": 63
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "IMAGE",
|
||||
"type": "IMAGE",
|
||||
"links": [
|
||||
64
|
||||
]
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "ImageScale"
|
||||
},
|
||||
"widgets_values": [
|
||||
"nearest-exact",
|
||||
512,
|
||||
512,
|
||||
"disabled"
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 33,
|
||||
"type": "ImageToMask",
|
||||
"pos": [
|
||||
21039.013671875,
|
||||
-430.8477783203125
|
||||
],
|
||||
"size": [
|
||||
270,
|
||||
58
|
||||
],
|
||||
"flags": {},
|
||||
"order": 8,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "image",
|
||||
"type": "IMAGE",
|
||||
"link": 67
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "MASK",
|
||||
"type": "MASK",
|
||||
"links": [
|
||||
68
|
||||
]
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "ImageToMask"
|
||||
},
|
||||
"widgets_values": [
|
||||
"red"
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 29,
|
||||
"type": "KV_Edit_Load",
|
||||
"pos": [
|
||||
20803.994140625,
|
||||
-1118.5433349609375
|
||||
],
|
||||
"size": [
|
||||
270,
|
||||
106
|
||||
],
|
||||
"flags": {},
|
||||
"order": 0,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "model",
|
||||
"shape": 7,
|
||||
"type": "MODEL",
|
||||
"link": null
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "model",
|
||||
"type": "MODEL_KVEDIT",
|
||||
"links": [
|
||||
69
|
||||
]
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "KV_Edit_Load"
|
||||
},
|
||||
"widgets_values": [
|
||||
"flux1-kj-dev-fp8.safetensors",
|
||||
true,
|
||||
true
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 34,
|
||||
"type": "KV_Edit_Load_LoRA",
|
||||
"pos": [
|
||||
21140.34765625,
|
||||
-1050.898193359375
|
||||
],
|
||||
"size": [
|
||||
270,
|
||||
82
|
||||
],
|
||||
"flags": {},
|
||||
"order": 4,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "model",
|
||||
"type": "MODEL_KVEDIT",
|
||||
"link": 69
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "model",
|
||||
"type": "MODEL_KVEDIT",
|
||||
"links": [
|
||||
70
|
||||
]
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "KV_Edit_Load_LoRA"
|
||||
},
|
||||
"widgets_values": [
|
||||
"flux-turbo.safetensors",
|
||||
1
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 5,
|
||||
"type": "LoadImage",
|
||||
"pos": [
|
||||
20210.099609375,
|
||||
-925.0836791992188
|
||||
],
|
||||
"size": [
|
||||
315,
|
||||
314
|
||||
],
|
||||
"flags": {},
|
||||
"order": 1,
|
||||
"mode": 0,
|
||||
"inputs": [],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "IMAGE",
|
||||
"type": "IMAGE",
|
||||
"slot_index": 0,
|
||||
"links": [
|
||||
63
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "MASK",
|
||||
"type": "MASK",
|
||||
"slot_index": 1,
|
||||
"links": [
|
||||
65
|
||||
]
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "LoadImage"
|
||||
},
|
||||
"widgets_values": [
|
||||
"clipspace/clipspace-painted-masked-1704959.png [input]",
|
||||
"image"
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 27,
|
||||
"type": "DualCLIPLoader",
|
||||
"pos": [
|
||||
20632.6328125,
|
||||
-937.8939208984375
|
||||
],
|
||||
"size": [
|
||||
315,
|
||||
130
|
||||
],
|
||||
"flags": {},
|
||||
"order": 2,
|
||||
"mode": 0,
|
||||
"inputs": [],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "CLIP",
|
||||
"type": "CLIP",
|
||||
"slot_index": 0,
|
||||
"links": [
|
||||
60
|
||||
]
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "DualCLIPLoader"
|
||||
},
|
||||
"widgets_values": [
|
||||
"t5xxl_fp8_e4m3fn.safetensors",
|
||||
"clip_l.safetensors",
|
||||
"flux",
|
||||
"default"
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 28,
|
||||
"type": "VAELoader",
|
||||
"pos": [
|
||||
20623.068359375,
|
||||
-738.1512451171875
|
||||
],
|
||||
"size": [
|
||||
315,
|
||||
58
|
||||
],
|
||||
"flags": {},
|
||||
"order": 3,
|
||||
"mode": 0,
|
||||
"inputs": [],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "VAE",
|
||||
"type": "VAE",
|
||||
"links": [
|
||||
61
|
||||
]
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "VAELoader"
|
||||
},
|
||||
"widgets_values": [
|
||||
"ae.safetensors"
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 25,
|
||||
"type": "KV_Edit_PreData",
|
||||
"pos": [
|
||||
20983.224609375,
|
||||
-812.3342895507812
|
||||
],
|
||||
"size": [
|
||||
400,
|
||||
202
|
||||
],
|
||||
"flags": {},
|
||||
"order": 9,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "clip",
|
||||
"type": "CLIP",
|
||||
"link": 60
|
||||
},
|
||||
{
|
||||
"name": "vae",
|
||||
"type": "VAE",
|
||||
"link": 61
|
||||
},
|
||||
{
|
||||
"name": "image",
|
||||
"type": "IMAGE",
|
||||
"link": 64
|
||||
},
|
||||
{
|
||||
"name": "mask",
|
||||
"type": "MASK",
|
||||
"link": 68
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "condition",
|
||||
"type": "CONDITION_KV",
|
||||
"links": [
|
||||
52
|
||||
]
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "KV_Edit_PreData"
|
||||
},
|
||||
"widgets_values": [
|
||||
"in a cluttered wooden cabin,a workbench holds a green neon sign that reads 'I love here'.",
|
||||
"in a cluttered wooden cabin,a workbench holds a green neon sign that reads 'I love nana'."
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 19,
|
||||
"type": "KV_Edit_Sampler",
|
||||
"pos": [
|
||||
21439.259765625,
|
||||
-876.9882202148438
|
||||
],
|
||||
"size": [
|
||||
315,
|
||||
294
|
||||
],
|
||||
"flags": {},
|
||||
"order": 10,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "model",
|
||||
"type": "MODEL_KVEDIT",
|
||||
"link": 70
|
||||
},
|
||||
{
|
||||
"name": "condition",
|
||||
"type": "CONDITION_KV",
|
||||
"link": 52
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "image",
|
||||
"type": "IMAGE",
|
||||
"slot_index": 0,
|
||||
"links": [
|
||||
46
|
||||
]
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "KV_Edit_Sampler"
|
||||
},
|
||||
"widgets_values": [
|
||||
968594714,
|
||||
"randomize",
|
||||
10,
|
||||
1.5,
|
||||
10,
|
||||
3.5,
|
||||
1,
|
||||
1,
|
||||
false,
|
||||
false
|
||||
]
|
||||
}
|
||||
],
|
||||
"links": [
|
||||
[
|
||||
46,
|
||||
19,
|
||||
0,
|
||||
14,
|
||||
0,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
52,
|
||||
25,
|
||||
0,
|
||||
19,
|
||||
1,
|
||||
"CONDITION_KV"
|
||||
],
|
||||
[
|
||||
60,
|
||||
27,
|
||||
0,
|
||||
25,
|
||||
0,
|
||||
"CLIP"
|
||||
],
|
||||
[
|
||||
61,
|
||||
28,
|
||||
0,
|
||||
25,
|
||||
1,
|
||||
"VAE"
|
||||
],
|
||||
[
|
||||
63,
|
||||
5,
|
||||
0,
|
||||
30,
|
||||
0,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
64,
|
||||
30,
|
||||
0,
|
||||
25,
|
||||
2,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
65,
|
||||
5,
|
||||
1,
|
||||
31,
|
||||
0,
|
||||
"MASK"
|
||||
],
|
||||
[
|
||||
66,
|
||||
31,
|
||||
0,
|
||||
32,
|
||||
0,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
67,
|
||||
32,
|
||||
0,
|
||||
33,
|
||||
0,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
68,
|
||||
33,
|
||||
0,
|
||||
25,
|
||||
3,
|
||||
"MASK"
|
||||
],
|
||||
[
|
||||
69,
|
||||
29,
|
||||
0,
|
||||
34,
|
||||
0,
|
||||
"MODEL_KVEDIT"
|
||||
],
|
||||
[
|
||||
70,
|
||||
34,
|
||||
0,
|
||||
19,
|
||||
0,
|
||||
"MODEL_KVEDIT"
|
||||
]
|
||||
],
|
||||
"groups": [],
|
||||
"config": {},
|
||||
"extra": {
|
||||
"ds": {
|
||||
"scale": 0.7400249944258334,
|
||||
"offset": [
|
||||
-19984.64368291623,
|
||||
1200.0752608333514
|
||||
]
|
||||
},
|
||||
"frontendVersion": "1.26.13",
|
||||
"VHS_latentpreview": false,
|
||||
"VHS_latentpreviewrate": 0,
|
||||
"VHS_MetadataImage": true,
|
||||
"VHS_KeepIntermediate": true
|
||||
},
|
||||
"version": 0.4
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,11 +0,0 @@
|
||||
try:
|
||||
from ._version import version as __version__ # type: ignore
|
||||
from ._version import version_tuple
|
||||
except ImportError:
|
||||
__version__ = "unknown (no version information available)"
|
||||
version_tuple = (0, 0, "unknown", "noinfo")
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
PACKAGE = __package__.replace("_", "-")
|
||||
PACKAGE_ROOT = Path(__file__).parent
|
||||
@@ -1,4 +0,0 @@
|
||||
# from .cli import app
|
||||
|
||||
# if __name__ == "__main__":
|
||||
# app()
|
||||
@@ -1,16 +0,0 @@
|
||||
# file generated by setuptools_scm
|
||||
# don't change, don't track in version control
|
||||
TYPE_CHECKING = False
|
||||
if TYPE_CHECKING:
|
||||
from typing import Tuple, Union
|
||||
VERSION_TUPLE = Tuple[Union[int, str], ...]
|
||||
else:
|
||||
VERSION_TUPLE = object
|
||||
|
||||
version: str
|
||||
__version__: str
|
||||
__version_tuple__: VERSION_TUPLE
|
||||
version_tuple: VERSION_TUPLE
|
||||
|
||||
__version__ = version = '0.0.post6+ge52e00f.d20250111'
|
||||
__version_tuple__ = version_tuple = (0, 0, 'ge52e00f.d20250111')
|
||||
-194
@@ -1,194 +0,0 @@
|
||||
import io
|
||||
import os
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import requests
|
||||
from PIL import Image
|
||||
|
||||
API_ENDPOINT = "https://api.bfl.ml"
|
||||
|
||||
|
||||
class ApiException(Exception):
|
||||
def __init__(self, status_code: int, detail: str | list[dict] | None = None):
|
||||
super().__init__()
|
||||
self.detail = detail
|
||||
self.status_code = status_code
|
||||
|
||||
def __str__(self) -> str:
|
||||
return self.__repr__()
|
||||
|
||||
def __repr__(self) -> str:
|
||||
if self.detail is None:
|
||||
message = None
|
||||
elif isinstance(self.detail, str):
|
||||
message = self.detail
|
||||
else:
|
||||
message = "[" + ",".join(d["msg"] for d in self.detail) + "]"
|
||||
return f"ApiException({self.status_code=}, {message=}, detail={self.detail})"
|
||||
|
||||
|
||||
class ImageRequest:
|
||||
def __init__(
|
||||
self,
|
||||
prompt: str,
|
||||
width: int = 1024,
|
||||
height: int = 1024,
|
||||
name: str = "flux.1-pro",
|
||||
num_steps: int = 50,
|
||||
prompt_upsampling: bool = False,
|
||||
seed: int | None = None,
|
||||
validate: bool = True,
|
||||
launch: bool = True,
|
||||
api_key: str | None = None,
|
||||
):
|
||||
"""
|
||||
Manages an image generation request to the API.
|
||||
|
||||
Args:
|
||||
prompt: Prompt to sample
|
||||
width: Width of the image in pixel
|
||||
height: Height of the image in pixel
|
||||
name: Name of the model
|
||||
num_steps: Number of network evaluations
|
||||
prompt_upsampling: Use prompt upsampling
|
||||
seed: Fix the generation seed
|
||||
validate: Run input validation
|
||||
launch: Directly launches request
|
||||
api_key: Your API key if not provided by the environment
|
||||
|
||||
Raises:
|
||||
ValueError: For invalid input
|
||||
ApiException: For errors raised from the API
|
||||
"""
|
||||
if validate:
|
||||
if name not in ["flux.1-pro"]:
|
||||
raise ValueError(f"Invalid model {name}")
|
||||
elif width % 32 != 0:
|
||||
raise ValueError(f"width must be divisible by 32, got {width}")
|
||||
elif not (256 <= width <= 1440):
|
||||
raise ValueError(f"width must be between 256 and 1440, got {width}")
|
||||
elif height % 32 != 0:
|
||||
raise ValueError(f"height must be divisible by 32, got {height}")
|
||||
elif not (256 <= height <= 1440):
|
||||
raise ValueError(f"height must be between 256 and 1440, got {height}")
|
||||
elif not (1 <= num_steps <= 50):
|
||||
raise ValueError(f"steps must be between 1 and 50, got {num_steps}")
|
||||
|
||||
self.request_json = {
|
||||
"prompt": prompt,
|
||||
"width": width,
|
||||
"height": height,
|
||||
"variant": name,
|
||||
"steps": num_steps,
|
||||
"prompt_upsampling": prompt_upsampling,
|
||||
}
|
||||
if seed is not None:
|
||||
self.request_json["seed"] = seed
|
||||
|
||||
self.request_id: str | None = None
|
||||
self.result: dict | None = None
|
||||
self._image_bytes: bytes | None = None
|
||||
self._url: str | None = None
|
||||
if api_key is None:
|
||||
self.api_key = os.environ.get("BFL_API_KEY")
|
||||
else:
|
||||
self.api_key = api_key
|
||||
|
||||
if launch:
|
||||
self.request()
|
||||
|
||||
def request(self):
|
||||
"""
|
||||
Request to generate the image.
|
||||
"""
|
||||
if self.request_id is not None:
|
||||
return
|
||||
response = requests.post(
|
||||
f"{API_ENDPOINT}/v1/image",
|
||||
headers={
|
||||
"accept": "application/json",
|
||||
"x-key": self.api_key,
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
json=self.request_json,
|
||||
)
|
||||
result = response.json()
|
||||
if response.status_code != 200:
|
||||
raise ApiException(status_code=response.status_code, detail=result.get("detail"))
|
||||
self.request_id = response.json()["id"]
|
||||
|
||||
def retrieve(self) -> dict:
|
||||
"""
|
||||
Wait for the generation to finish and retrieve response.
|
||||
"""
|
||||
if self.request_id is None:
|
||||
self.request()
|
||||
while self.result is None:
|
||||
response = requests.get(
|
||||
f"{API_ENDPOINT}/v1/get_result",
|
||||
headers={
|
||||
"accept": "application/json",
|
||||
"x-key": self.api_key,
|
||||
},
|
||||
params={
|
||||
"id": self.request_id,
|
||||
},
|
||||
)
|
||||
result = response.json()
|
||||
if "status" not in result:
|
||||
raise ApiException(status_code=response.status_code, detail=result.get("detail"))
|
||||
elif result["status"] == "Ready":
|
||||
self.result = result["result"]
|
||||
elif result["status"] == "Pending":
|
||||
time.sleep(0.5)
|
||||
else:
|
||||
raise ApiException(status_code=200, detail=f"API returned status '{result['status']}'")
|
||||
return self.result
|
||||
|
||||
@property
|
||||
def bytes(self) -> bytes:
|
||||
"""
|
||||
Generated image as bytes.
|
||||
"""
|
||||
if self._image_bytes is None:
|
||||
response = requests.get(self.url)
|
||||
if response.status_code == 200:
|
||||
self._image_bytes = response.content
|
||||
else:
|
||||
raise ApiException(status_code=response.status_code)
|
||||
return self._image_bytes
|
||||
|
||||
@property
|
||||
def url(self) -> str:
|
||||
"""
|
||||
Public url to retrieve the image from
|
||||
"""
|
||||
if self._url is None:
|
||||
result = self.retrieve()
|
||||
self._url = result["sample"]
|
||||
return self._url
|
||||
|
||||
@property
|
||||
def image(self) -> Image.Image:
|
||||
"""
|
||||
Load the image as a PIL Image
|
||||
"""
|
||||
return Image.open(io.BytesIO(self.bytes))
|
||||
|
||||
def save(self, path: str):
|
||||
"""
|
||||
Save the generated image to a local path
|
||||
"""
|
||||
suffix = Path(self.url).suffix
|
||||
if not path.endswith(suffix):
|
||||
path = path + suffix
|
||||
Path(path).resolve().parent.mkdir(parents=True, exist_ok=True)
|
||||
with open(path, "wb") as file:
|
||||
file.write(self.bytes)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
from fire import Fire
|
||||
|
||||
Fire(ImageRequest)
|
||||
@@ -1,19 +0,0 @@
|
||||
{
|
||||
"_class_name": "FluxTransformer2DModel",
|
||||
"_diffusers_version": "0.30.0.dev0",
|
||||
"_name_or_path": "../checkpoints/flux-dev/transformer",
|
||||
"attention_head_dim": 128,
|
||||
"axes_dims_rope": [
|
||||
16,
|
||||
56,
|
||||
56
|
||||
],
|
||||
"guidance_embeds": true,
|
||||
"in_channels": 64,
|
||||
"joint_attention_dim": 4096,
|
||||
"num_attention_heads": 24,
|
||||
"num_layers": 19,
|
||||
"num_single_layers": 38,
|
||||
"patch_size": 1,
|
||||
"pooled_projection_dim": 768
|
||||
}
|
||||
-248
@@ -1,248 +0,0 @@
|
||||
from dataclasses import dataclass
|
||||
from .util import load_flux_model_,load_flux_model
|
||||
from einops import rearrange,repeat
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import Tensor
|
||||
|
||||
from .sampling import get_schedule, unpack,denoise_kv,denoise_kv_inf
|
||||
from .util import load_flux_model_,load_flux_model
|
||||
from .model import Flux_kv
|
||||
|
||||
# @dataclass
|
||||
# class SamplingOptions:
|
||||
# source_prompt: str = ''
|
||||
# target_prompt: str = ''
|
||||
# # prompt: str
|
||||
# width: int = 1366
|
||||
# height: int = 768
|
||||
# inversion_num_steps: int = 0
|
||||
# denoise_num_steps: int = 0
|
||||
# skip_step: int = 0
|
||||
# inversion_guidance: float = 1.0
|
||||
# denoise_guidance: float = 1.0
|
||||
# seed: int = 42
|
||||
# re_init: bool = False
|
||||
# attn_mask: bool = False
|
||||
|
||||
class only_Flux(torch.nn.Module):
|
||||
def __init__(self, device,model,name='flux-dev'):
|
||||
super().__init__()
|
||||
self.device = device
|
||||
self.name = name
|
||||
self.is_dev = self.name != "flux-schnell"
|
||||
#self.model = load_flow_model(self.name, device=self.device,flux_cls=Flux_kv)
|
||||
if isinstance(model, str):
|
||||
self.model = load_flux_model(model,self.device, flux_cls=Flux_kv)
|
||||
else:
|
||||
self.model = load_flux_model_(model,self.device, flux_cls=Flux_kv)
|
||||
|
||||
def create_attention_mask(self,seq_len, mask_indices, text_len=512 , device='cuda'):
|
||||
# """
|
||||
# 创建自定义的注意力掩码。
|
||||
|
||||
# Args:
|
||||
# seq_len (int): 序列长度。
|
||||
# mask_indices (List[int]): 图像令牌中掩码区域的索引。
|
||||
# text_len (int): 文本令牌的长度，默认 512。改成comfy的时候用256
|
||||
# device (str): 设备类型，如 'cuda' 或 'cpu'。
|
||||
|
||||
# Returns:
|
||||
# torch.Tensor: 形状为 (seq_len, seq_len) 的注意力掩码。
|
||||
# """
|
||||
# 初始化掩码为全 False
|
||||
|
||||
|
||||
#text_len=512 if self.is_dev else 256
|
||||
|
||||
attention_mask = torch.zeros(seq_len, seq_len, dtype=torch.bool, device=device)
|
||||
|
||||
# 文本令牌索引
|
||||
text_indices = torch.arange(0, text_len, device=device)
|
||||
|
||||
# 掩码区域令牌索引
|
||||
mask_token_indices = torch.tensor([idx + text_len for idx in mask_indices], device=device)
|
||||
|
||||
# 背景区域令牌索引
|
||||
all_indices = torch.arange(text_len, seq_len, device=device)
|
||||
background_token_indices = torch.tensor([idx for idx in all_indices if idx not in mask_token_indices])
|
||||
|
||||
# 设置文本查询可以关注所有键
|
||||
attention_mask[text_indices.unsqueeze(1).expand(-1, seq_len)] = True
|
||||
attention_mask[text_indices.unsqueeze(1), text_indices] = True# 关注文本
|
||||
attention_mask[text_indices.unsqueeze(1), background_token_indices] = True # 关注背景
|
||||
|
||||
|
||||
# attention_mask[mask_token_indices.unsqueeze(1), background_token_indices] = True # 关注背景
|
||||
attention_mask[mask_token_indices.unsqueeze(1), text_indices] = True # 关注文本
|
||||
attention_mask[mask_token_indices.unsqueeze(1), mask_token_indices] = True # 关注掩码区域
|
||||
|
||||
|
||||
# attention_mask[background_token_indices.unsqueeze(1).expand(-1, seq_len), :] = False
|
||||
# attention_mask[background_token_indices.unsqueeze(1), mask_token_indices] = True # 关注掩码
|
||||
attention_mask[background_token_indices.unsqueeze(1), text_indices] = True # 关注文本
|
||||
attention_mask[background_token_indices.unsqueeze(1), background_token_indices] = True # 关注背景区域
|
||||
|
||||
return attention_mask.unsqueeze(0)
|
||||
def create_attention_scale(self,seq_len, mask_indices,text_len=512, device='cuda',scale = 0):
|
||||
|
||||
attention_scale = torch.zeros(1, seq_len, dtype=torch.bfloat16, device=device) # 相加时广播
|
||||
|
||||
|
||||
text_indices = torch.arange(0, text_len, device=device)
|
||||
|
||||
mask_token_indices = torch.tensor([idx + text_len for idx in mask_indices], device=device)
|
||||
|
||||
all_indices = torch.arange(text_len, seq_len, device=device)
|
||||
background_token_indices = torch.tensor([idx for idx in all_indices if idx not in mask_token_indices])
|
||||
|
||||
attention_scale[0, background_token_indices] = scale #
|
||||
|
||||
return attention_scale.unsqueeze(0)
|
||||
|
||||
class Flux_kv_edit_inf(only_Flux):
|
||||
def __init__(self, device,model,name):
|
||||
super().__init__(device,model,name)
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def forward(self,inp,inp_target,mask,opts):
|
||||
|
||||
info = {}
|
||||
info['feature'] = {}
|
||||
bs, L, d = inp["img"].shape
|
||||
|
||||
_,tokens_L,_= inp["txt"].shape #comfy tokens_L=256
|
||||
|
||||
h = opts.height // 8
|
||||
w = opts.width // 8
|
||||
mask = F.interpolate(mask, size=(h,w), mode='bilinear', align_corners=False)
|
||||
mask[mask > 0] = 1
|
||||
#print("mask",mask.shape) #mask torch.Size([1, 1, 64, 64])
|
||||
mask = repeat(mask, 'b c h w -> b (repeat c) h w', repeat=16)
|
||||
mask = rearrange(mask, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2)
|
||||
info['mask'] = mask
|
||||
bool_mask = (mask.sum(dim=2) > 0.5)
|
||||
info['mask_indices'] = torch.nonzero(bool_mask)[:,1]
|
||||
#单独分离inversion
|
||||
if opts.attn_mask and (~bool_mask).any():
|
||||
attention_mask = self.create_attention_mask(L+tokens_L, info['mask_indices'],tokens_L, device=self.device) #使用512令牌时须将256改成512
|
||||
print("attention_mask",attention_mask.shape)
|
||||
else:
|
||||
attention_mask = None
|
||||
info['attention_mask'] = attention_mask
|
||||
|
||||
if opts.attn_scale != 0 and (~bool_mask).any():
|
||||
attention_scale = self.create_attention_scale(L+tokens_L, info['mask_indices'], tokens_L,device=mask.device,scale = opts.attn_scale)#使用512令牌时须将256改成512
|
||||
else:
|
||||
attention_scale = None
|
||||
info['attention_scale'] = attention_scale
|
||||
|
||||
denoise_timesteps = get_schedule(opts.denoise_num_steps, inp["img"].shape[1], shift=(self.name != "flux-schnell"))
|
||||
denoise_timesteps = denoise_timesteps[opts.skip_step:]
|
||||
|
||||
z0 = inp["img"]
|
||||
|
||||
with torch.no_grad():
|
||||
info['inverse'] = True
|
||||
z_fe, info = denoise_kv_inf(self.model, img=inp["img"], img_ids=inp['img_ids'],
|
||||
source_txt=inp['txt'], source_txt_ids=inp['txt_ids'], source_vec=inp['vec'],
|
||||
target_txt=inp_target['txt'], target_txt_ids=inp_target['txt_ids'], target_vec=inp_target['vec'],
|
||||
timesteps=denoise_timesteps, source_guidance=opts.inversion_guidance, target_guidance=opts.denoise_guidance,
|
||||
info=info)
|
||||
mask_indices = info['mask_indices']
|
||||
|
||||
z0[:, mask_indices,...] = z_fe
|
||||
|
||||
z0 = unpack(z0.float(), opts.height, opts.width)
|
||||
del info
|
||||
return z0
|
||||
|
||||
class Flux_kv_edit(only_Flux):
|
||||
def __init__( self, device,model_path,name):
|
||||
super().__init__(device,model_path,name)
|
||||
|
||||
@torch.inference_mode()
|
||||
def forward(self,inp,inp_target,mask:Tensor,opts):
|
||||
z0,zt,info = self.inverse(inp,mask,opts)
|
||||
z0 = self.denoise(z0,zt,inp_target,mask,opts,info)
|
||||
return z0
|
||||
@torch.inference_mode()
|
||||
def inverse(self,inp,mask,opts):
|
||||
info = {}
|
||||
info['feature'] = {}
|
||||
bs, L, d = inp["img"].shape
|
||||
_,tokens_L,_= inp["txt"].shape #comfy tokens_L=256
|
||||
h = opts.height // 8
|
||||
w = opts.width // 8
|
||||
|
||||
if opts.attn_mask:
|
||||
mask = F.interpolate(mask, size=(h,w), mode='bilinear', align_corners=False)
|
||||
mask[mask > 0] = 1
|
||||
|
||||
mask = repeat(mask, 'b c h w -> b (repeat c) h w', repeat=16)
|
||||
mask = rearrange(mask, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2)
|
||||
bool_mask = (mask.sum(dim=2) > 0.5)
|
||||
mask_indices = torch.nonzero(bool_mask)[:,1]
|
||||
|
||||
#单独分离inversion
|
||||
assert not (~bool_mask).all(), "mask is all false"
|
||||
assert not (bool_mask).all(), "mask is all true"
|
||||
attention_mask = self.create_attention_mask(L+tokens_L, mask_indices,tokens_L, device=mask.device)#512->256
|
||||
info['attention_mask'] = attention_mask
|
||||
|
||||
|
||||
denoise_timesteps = get_schedule(opts.denoise_num_steps, inp["img"].shape[1], shift=(self.name != "flux-schnell"))
|
||||
denoise_timesteps = denoise_timesteps[opts.skip_step:]
|
||||
|
||||
# 加噪过程
|
||||
z0 = inp["img"].clone()
|
||||
info['inverse'] = True
|
||||
zt, info = denoise_kv(self.model, **inp, timesteps=denoise_timesteps, guidance=opts.inversion_guidance, inverse=True, info=info)
|
||||
return z0,zt,info
|
||||
|
||||
@torch.inference_mode()
|
||||
def denoise(self,z0,zt,inp_target,mask:Tensor,opts,info):
|
||||
|
||||
h = opts.height // 8
|
||||
w = opts.width // 8
|
||||
L = h * w // 4
|
||||
_,tokens_L,_= inp_target["txt"].shape #comfy tokens_L=256
|
||||
mask = F.interpolate(mask, size=(h,w), mode='bilinear', align_corners=False)
|
||||
mask[mask > 0] = 1
|
||||
|
||||
mask = repeat(mask, 'b c h w -> b (repeat c) h w', repeat=16)
|
||||
|
||||
mask = rearrange(mask, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2)
|
||||
info['mask'] = mask
|
||||
bool_mask = (mask.sum(dim=2) > 0.5)
|
||||
info['mask_indices'] = torch.nonzero(bool_mask)[:,1]
|
||||
|
||||
denoise_timesteps = get_schedule(opts.denoise_num_steps, inp_target["img"].shape[1], shift=(self.name != "flux-schnell"))
|
||||
denoise_timesteps = denoise_timesteps[opts.skip_step:]
|
||||
|
||||
mask_indices = info['mask_indices']
|
||||
if opts.re_init:
|
||||
noise = torch.randn_like(zt)
|
||||
t = denoise_timesteps[0]
|
||||
zt_noise = z0 *(1 - t) + noise * t
|
||||
inp_target["img"] = zt_noise[:, mask_indices,...]
|
||||
else:
|
||||
img_name = str(info['t']) + '_' + 'img'
|
||||
zt = info['feature'][img_name].to(zt.device)
|
||||
inp_target["img"] = zt[:, mask_indices,...]
|
||||
|
||||
if opts.attn_scale != 0 and (~bool_mask).any():
|
||||
attention_scale = self.create_attention_scale(L+tokens_L, mask_indices,tokens_L, device=mask.device,scale = opts.attn_scale)#512->256
|
||||
else:
|
||||
attention_scale = None
|
||||
info['attention_scale'] = attention_scale
|
||||
|
||||
info['inverse'] = False
|
||||
x, _ = denoise_kv(self.model, **inp_target, timesteps=denoise_timesteps, guidance=opts.denoise_guidance, inverse=False, info=info)
|
||||
|
||||
z0[:, mask_indices,...] = z0[:, mask_indices,...] * (1 - info['mask'][:, mask_indices,...]) + x * info['mask'][:, mask_indices,...]
|
||||
|
||||
z0 = unpack(z0.float(), opts.height, opts.width)
|
||||
del info
|
||||
return z0
|
||||
@@ -1,40 +0,0 @@
|
||||
import torch
|
||||
from einops import rearrange
|
||||
from torch import Tensor
|
||||
|
||||
|
||||
def attention(q: Tensor, k: Tensor, v: Tensor, pe: Tensor,pe_q = None, attention_mask = None) -> Tensor:
|
||||
if pe_q is None:
|
||||
q, k = apply_rope(q, k, pe)
|
||||
x = torch.nn.functional.scaled_dot_product_attention(q, k, v,attn_mask=attention_mask)
|
||||
x = rearrange(x, "B H L D -> B L (H D)")
|
||||
return x
|
||||
else:
|
||||
q, k = apply_rope_qk(q, k, pe_q, pe)
|
||||
x = torch.nn.functional.scaled_dot_product_attention(q, k, v,attn_mask=attention_mask)
|
||||
x = rearrange(x, "B H L D -> B L (H D)")
|
||||
return x
|
||||
|
||||
def rope(pos: Tensor, dim: int, theta: int) -> Tensor:
|
||||
assert dim % 2 == 0
|
||||
scale = torch.arange(0, dim, 2, dtype=torch.float64, device=pos.device) / dim # dim =16 + 56 + 56
|
||||
omega = 1.0 / (theta**scale) # 64 omega
|
||||
out = torch.einsum("...n,d->...nd", pos, omega)
|
||||
out = torch.stack([torch.cos(out), -torch.sin(out), torch.sin(out), torch.cos(out)], dim=-1)
|
||||
out = rearrange(out, "b n d (i j) -> b n d i j", i=2, j=2)
|
||||
return out.float()
|
||||
|
||||
|
||||
def apply_rope(xq: Tensor, xk: Tensor, freqs_cis: Tensor) -> tuple[Tensor, Tensor]:
|
||||
xq_ = xq.float().reshape(*xq.shape[:-1], -1, 1, 2)
|
||||
xk_ = xk.float().reshape(*xk.shape[:-1], -1, 1, 2)
|
||||
xq_out = freqs_cis[:, :, :xq_.shape[2], :, :, 0] * xq_[..., 0] + freqs_cis[:, :, :xq_.shape[2], :, :, 1] * xq_[..., 1]
|
||||
xk_out = freqs_cis[:, :, :xk_.shape[2], :, :, 0] * xk_[..., 0] + freqs_cis[:, :, :xk_.shape[2], :, :, 1] * xk_[..., 1]
|
||||
return xq_out.reshape(*xq.shape).type_as(xq), xk_out.reshape(*xk.shape).type_as(xk)
|
||||
|
||||
def apply_rope_qk(xq: Tensor, xk: Tensor, freqs_cis_q: Tensor,freqs_cis_k: Tensor) -> tuple[Tensor, Tensor]:
|
||||
xq_ = xq.float().reshape(*xq.shape[:-1], -1, 1, 2)
|
||||
xk_ = xk.float().reshape(*xk.shape[:-1], -1, 1, 2)
|
||||
xq_out = freqs_cis_q[:, :, :xq_.shape[2], :, :, 0] * xq_[..., 0] + freqs_cis_q[:, :, :xq_.shape[2], :, :, 1] * xq_[..., 1]
|
||||
xk_out = freqs_cis_k[:, :, :xk_.shape[2], :, :, 0] * xk_[..., 0] + freqs_cis_k[:, :, :xk_.shape[2], :, :, 1] * xk_[..., 1]
|
||||
return xq_out.reshape(*xq.shape).type_as(xq), xk_out.reshape(*xk.shape).type_as(xk)
|
||||
-172
@@ -1,172 +0,0 @@
|
||||
from dataclasses import dataclass
|
||||
|
||||
import torch
|
||||
from torch import Tensor, nn
|
||||
|
||||
from .modules.layers import (DoubleStreamBlock, EmbedND, LastLayer,
|
||||
MLPEmbedder, SingleStreamBlock,
|
||||
SingleStreamBlock_kv,DoubleStreamBlock_kv,
|
||||
timestep_embedding)
|
||||
|
||||
|
||||
@dataclass
|
||||
class FluxParams:
|
||||
in_channels: int
|
||||
vec_in_dim: int
|
||||
context_in_dim: int
|
||||
hidden_size: int
|
||||
mlp_ratio: float
|
||||
num_heads: int
|
||||
depth: int
|
||||
depth_single_blocks: int
|
||||
axes_dim: list[int]
|
||||
theta: int
|
||||
qkv_bias: bool
|
||||
guidance_embed: bool
|
||||
|
||||
|
||||
class Flux(nn.Module):
|
||||
"""
|
||||
Transformer model for flow matching on sequences.
|
||||
"""
|
||||
|
||||
def __init__(self, params: FluxParams,double_block_cls=DoubleStreamBlock,single_block_cls=SingleStreamBlock):
|
||||
super().__init__()
|
||||
|
||||
self.params = params
|
||||
self.in_channels = params.in_channels
|
||||
self.out_channels = self.in_channels
|
||||
if params.hidden_size % params.num_heads != 0:
|
||||
raise ValueError(
|
||||
f"Hidden size {params.hidden_size} must be divisible by num_heads {params.num_heads}"
|
||||
)
|
||||
pe_dim = params.hidden_size // params.num_heads
|
||||
if sum(params.axes_dim) != pe_dim:
|
||||
raise ValueError(f"Got {params.axes_dim} but expected positional dim {pe_dim}")
|
||||
self.hidden_size = params.hidden_size
|
||||
self.num_heads = params.num_heads
|
||||
self.pe_embedder = EmbedND(dim=pe_dim, theta=params.theta, axes_dim=params.axes_dim)
|
||||
self.img_in = nn.Linear(self.in_channels, self.hidden_size, bias=True)
|
||||
self.time_in = MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size)
|
||||
self.vector_in = MLPEmbedder(params.vec_in_dim, self.hidden_size)
|
||||
self.guidance_in = (
|
||||
MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size) if params.guidance_embed else nn.Identity()
|
||||
)
|
||||
self.txt_in = nn.Linear(params.context_in_dim, self.hidden_size)
|
||||
|
||||
self.double_blocks = nn.ModuleList(
|
||||
[
|
||||
double_block_cls(
|
||||
self.hidden_size,
|
||||
self.num_heads,
|
||||
mlp_ratio=params.mlp_ratio,
|
||||
qkv_bias=params.qkv_bias,
|
||||
)
|
||||
for _ in range(params.depth)
|
||||
]
|
||||
)
|
||||
|
||||
self.single_blocks = nn.ModuleList(
|
||||
[
|
||||
single_block_cls(self.hidden_size, self.num_heads, mlp_ratio=params.mlp_ratio)
|
||||
for _ in range(params.depth_single_blocks)
|
||||
]
|
||||
)
|
||||
|
||||
self.final_layer = LastLayer(self.hidden_size, 1, self.out_channels)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
img: Tensor,
|
||||
img_ids: Tensor,
|
||||
txt: Tensor,
|
||||
txt_ids: Tensor,
|
||||
timesteps: Tensor,
|
||||
y: Tensor,
|
||||
guidance: Tensor | None = 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)
|
||||
vec = self.time_in(timestep_embedding(timesteps, 256))
|
||||
if self.params.guidance_embed:
|
||||
if guidance is None:
|
||||
raise ValueError("Didn't get guidance strength for guidance distilled model.")
|
||||
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)
|
||||
|
||||
for block in self.double_blocks:
|
||||
img, txt = block(img=img, txt=txt, vec=vec, pe=pe)
|
||||
|
||||
img = torch.cat((txt, img), 1)
|
||||
for block in self.single_blocks:
|
||||
img = block(img, vec=vec, pe=pe)
|
||||
img = img[:, txt.shape[1] :, ...]
|
||||
|
||||
img = self.final_layer(img, vec) # (N, T, patch_size ** 2 * out_channels)
|
||||
return img
|
||||
|
||||
class Flux_kv(Flux):
|
||||
"""
|
||||
继承Flux类，重写forward方法
|
||||
"""
|
||||
|
||||
def __init__(self, params: FluxParams,double_block_cls=DoubleStreamBlock_kv,single_block_cls=SingleStreamBlock_kv):
|
||||
super().__init__(params,double_block_cls,single_block_cls)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
img: Tensor,
|
||||
img_ids: Tensor,
|
||||
txt: Tensor,
|
||||
txt_ids: Tensor,
|
||||
timesteps: Tensor,
|
||||
y: Tensor,
|
||||
guidance: Tensor | None = None,
|
||||
info: dict = {},
|
||||
) -> Tensor:
|
||||
if img.ndim != 3 or txt.ndim != 3:
|
||||
raise ValueError("Input img and txt tensors must have 3 dimensions.")
|
||||
|
||||
token_len=txt.shape[1]
|
||||
# running on sequences img
|
||||
img = self.img_in(img)
|
||||
vec = self.time_in(timestep_embedding(timesteps, 256))
|
||||
if self.params.guidance_embed:
|
||||
if guidance is None:
|
||||
raise ValueError("Didn't get guidance strength for guidance distilled model.")
|
||||
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)
|
||||
if not info['inverse']:
|
||||
mask_indices = info['mask_indices']
|
||||
info['pe_mask'] = torch.cat((pe[:, :, :token_len, ...],pe[:, :, mask_indices+token_len, ...]),dim=2)
|
||||
|
||||
cnt = 0
|
||||
for block in self.double_blocks:
|
||||
info['id'] = cnt
|
||||
img, txt = block(img=img, txt=txt, vec=vec, pe=pe, info=info)
|
||||
cnt += 1
|
||||
|
||||
cnt = 0
|
||||
x = torch.cat((txt, img), 1)
|
||||
info["token_len"] = token_len
|
||||
for block in self.single_blocks:
|
||||
info['id'] = cnt
|
||||
x = block(x, vec=vec, pe=pe, info=info)
|
||||
cnt += 1
|
||||
|
||||
img = x[:, txt.shape[1] :, ...]
|
||||
|
||||
img = self.final_layer(img, vec) # (N, T, patch_size ** 2 * out_channels)
|
||||
|
||||
return img
|
||||
@@ -1,313 +0,0 @@
|
||||
from dataclasses import dataclass
|
||||
|
||||
import torch
|
||||
from einops import rearrange
|
||||
from torch import Tensor, nn
|
||||
|
||||
|
||||
@dataclass
|
||||
class AutoEncoderParams:
|
||||
resolution: int
|
||||
in_channels: int
|
||||
ch: int
|
||||
out_ch: int
|
||||
ch_mult: list[int]
|
||||
num_res_blocks: int
|
||||
z_channels: int
|
||||
scale_factor: float
|
||||
shift_factor: float
|
||||
|
||||
|
||||
def swish(x: Tensor) -> Tensor:
|
||||
return x * torch.sigmoid(x)
|
||||
|
||||
|
||||
class AttnBlock(nn.Module):
|
||||
def __init__(self, in_channels: int):
|
||||
super().__init__()
|
||||
self.in_channels = in_channels
|
||||
|
||||
self.norm = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
|
||||
|
||||
self.q = nn.Conv2d(in_channels, in_channels, kernel_size=1)
|
||||
self.k = nn.Conv2d(in_channels, in_channels, kernel_size=1)
|
||||
self.v = nn.Conv2d(in_channels, in_channels, kernel_size=1)
|
||||
self.proj_out = nn.Conv2d(in_channels, in_channels, kernel_size=1)
|
||||
|
||||
def attention(self, h_: Tensor) -> Tensor:
|
||||
h_ = self.norm(h_)
|
||||
q = self.q(h_)
|
||||
k = self.k(h_)
|
||||
v = self.v(h_)
|
||||
|
||||
b, c, h, w = q.shape
|
||||
q = rearrange(q, "b c h w -> b 1 (h w) c").contiguous()
|
||||
k = rearrange(k, "b c h w -> b 1 (h w) c").contiguous()
|
||||
v = rearrange(v, "b c h w -> b 1 (h w) c").contiguous()
|
||||
h_ = nn.functional.scaled_dot_product_attention(q, k, v)
|
||||
|
||||
return rearrange(h_, "b 1 (h w) c -> b c h w", h=h, w=w, c=c, b=b)
|
||||
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
return x + self.proj_out(self.attention(x))
|
||||
|
||||
|
||||
class ResnetBlock(nn.Module):
|
||||
def __init__(self, in_channels: int, out_channels: int):
|
||||
super().__init__()
|
||||
self.in_channels = in_channels
|
||||
out_channels = in_channels if out_channels is None else out_channels
|
||||
self.out_channels = out_channels
|
||||
|
||||
self.norm1 = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
|
||||
self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
|
||||
self.norm2 = nn.GroupNorm(num_groups=32, num_channels=out_channels, eps=1e-6, affine=True)
|
||||
self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
|
||||
if self.in_channels != self.out_channels:
|
||||
self.nin_shortcut = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
|
||||
|
||||
def forward(self, x):
|
||||
h = x
|
||||
h = self.norm1(h)
|
||||
h = swish(h)
|
||||
h = self.conv1(h)
|
||||
|
||||
h = self.norm2(h)
|
||||
h = swish(h)
|
||||
h = self.conv2(h)
|
||||
|
||||
if self.in_channels != self.out_channels:
|
||||
x = self.nin_shortcut(x)
|
||||
|
||||
return x + h
|
||||
|
||||
|
||||
class Downsample(nn.Module):
|
||||
def __init__(self, in_channels: int):
|
||||
super().__init__()
|
||||
# no asymmetric padding in torch conv, must do it ourselves
|
||||
self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=2, padding=0)
|
||||
|
||||
def forward(self, x: Tensor):
|
||||
pad = (0, 1, 0, 1)
|
||||
x = nn.functional.pad(x, pad, mode="constant", value=0)
|
||||
x = self.conv(x)
|
||||
return x
|
||||
|
||||
|
||||
class Upsample(nn.Module):
|
||||
def __init__(self, in_channels: int):
|
||||
super().__init__()
|
||||
self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1)
|
||||
|
||||
def forward(self, x: Tensor):
|
||||
x = nn.functional.interpolate(x, scale_factor=2.0, mode="nearest")
|
||||
x = self.conv(x)
|
||||
return x
|
||||
|
||||
|
||||
class Encoder(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
resolution: int,
|
||||
in_channels: int,
|
||||
ch: int,
|
||||
ch_mult: list[int],
|
||||
num_res_blocks: int,
|
||||
z_channels: int,
|
||||
):
|
||||
super().__init__()
|
||||
self.ch = ch
|
||||
self.num_resolutions = len(ch_mult)
|
||||
self.num_res_blocks = num_res_blocks
|
||||
self.resolution = resolution
|
||||
self.in_channels = in_channels
|
||||
# downsampling
|
||||
self.conv_in = nn.Conv2d(in_channels, self.ch, kernel_size=3, stride=1, padding=1)
|
||||
|
||||
curr_res = resolution
|
||||
in_ch_mult = (1,) + tuple(ch_mult)
|
||||
self.in_ch_mult = in_ch_mult
|
||||
self.down = nn.ModuleList()
|
||||
block_in = self.ch
|
||||
for i_level in range(self.num_resolutions):
|
||||
block = nn.ModuleList()
|
||||
attn = nn.ModuleList()
|
||||
block_in = ch * in_ch_mult[i_level]
|
||||
block_out = ch * ch_mult[i_level]
|
||||
for _ in range(self.num_res_blocks):
|
||||
block.append(ResnetBlock(in_channels=block_in, out_channels=block_out))
|
||||
block_in = block_out
|
||||
down = nn.Module()
|
||||
down.block = block
|
||||
down.attn = attn
|
||||
if i_level != self.num_resolutions - 1:
|
||||
down.downsample = Downsample(block_in)
|
||||
curr_res = curr_res // 2
|
||||
self.down.append(down)
|
||||
|
||||
# middle
|
||||
self.mid = nn.Module()
|
||||
self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in)
|
||||
self.mid.attn_1 = AttnBlock(block_in)
|
||||
self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in)
|
||||
|
||||
# end
|
||||
self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in, eps=1e-6, affine=True)
|
||||
self.conv_out = nn.Conv2d(block_in, 2 * z_channels, kernel_size=3, stride=1, padding=1)
|
||||
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
# downsampling
|
||||
hs = [self.conv_in(x)]
|
||||
for i_level in range(self.num_resolutions):
|
||||
for i_block in range(self.num_res_blocks):
|
||||
h = self.down[i_level].block[i_block](hs[-1])
|
||||
if len(self.down[i_level].attn) > 0:
|
||||
h = self.down[i_level].attn[i_block](h)
|
||||
hs.append(h)
|
||||
if i_level != self.num_resolutions - 1:
|
||||
hs.append(self.down[i_level].downsample(hs[-1]))
|
||||
|
||||
# middle
|
||||
h = hs[-1]
|
||||
h = self.mid.block_1(h)
|
||||
h = self.mid.attn_1(h)
|
||||
h = self.mid.block_2(h)
|
||||
# end
|
||||
h = self.norm_out(h)
|
||||
h = swish(h)
|
||||
h = self.conv_out(h)
|
||||
return h
|
||||
|
||||
|
||||
class Decoder(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
ch: int,
|
||||
out_ch: int,
|
||||
ch_mult: list[int],
|
||||
num_res_blocks: int,
|
||||
in_channels: int,
|
||||
resolution: int,
|
||||
z_channels: int,
|
||||
):
|
||||
super().__init__()
|
||||
self.ch = ch
|
||||
self.num_resolutions = len(ch_mult)
|
||||
self.num_res_blocks = num_res_blocks
|
||||
self.resolution = resolution
|
||||
self.in_channels = in_channels
|
||||
self.ffactor = 2 ** (self.num_resolutions - 1)
|
||||
|
||||
# compute in_ch_mult, block_in and curr_res at lowest res
|
||||
block_in = ch * ch_mult[self.num_resolutions - 1]
|
||||
curr_res = resolution // 2 ** (self.num_resolutions - 1)
|
||||
self.z_shape = (1, z_channels, curr_res, curr_res)
|
||||
|
||||
# z to block_in
|
||||
self.conv_in = nn.Conv2d(z_channels, block_in, kernel_size=3, stride=1, padding=1)
|
||||
|
||||
# middle
|
||||
self.mid = nn.Module()
|
||||
self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in)
|
||||
self.mid.attn_1 = AttnBlock(block_in)
|
||||
self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in)
|
||||
|
||||
# upsampling
|
||||
self.up = nn.ModuleList()
|
||||
for i_level in reversed(range(self.num_resolutions)):
|
||||
block = nn.ModuleList()
|
||||
attn = nn.ModuleList()
|
||||
block_out = ch * ch_mult[i_level]
|
||||
for _ in range(self.num_res_blocks + 1):
|
||||
block.append(ResnetBlock(in_channels=block_in, out_channels=block_out))
|
||||
block_in = block_out
|
||||
up = nn.Module()
|
||||
up.block = block
|
||||
up.attn = attn
|
||||
if i_level != 0:
|
||||
up.upsample = Upsample(block_in)
|
||||
curr_res = curr_res * 2
|
||||
self.up.insert(0, up) # prepend to get consistent order
|
||||
|
||||
# end
|
||||
self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in, eps=1e-6, affine=True)
|
||||
self.conv_out = nn.Conv2d(block_in, out_ch, kernel_size=3, stride=1, padding=1)
|
||||
|
||||
def forward(self, z: Tensor) -> Tensor:
|
||||
# z to block_in
|
||||
h = self.conv_in(z)
|
||||
|
||||
# middle
|
||||
h = self.mid.block_1(h)
|
||||
h = self.mid.attn_1(h)
|
||||
h = self.mid.block_2(h)
|
||||
|
||||
# upsampling
|
||||
for i_level in reversed(range(self.num_resolutions)):
|
||||
for i_block in range(self.num_res_blocks + 1):
|
||||
h = self.up[i_level].block[i_block](h)
|
||||
if len(self.up[i_level].attn) > 0:
|
||||
h = self.up[i_level].attn[i_block](h)
|
||||
if i_level != 0:
|
||||
h = self.up[i_level].upsample(h)
|
||||
|
||||
# end
|
||||
h = self.norm_out(h)
|
||||
h = swish(h)
|
||||
h = self.conv_out(h)
|
||||
return h
|
||||
|
||||
|
||||
class DiagonalGaussian(nn.Module):
|
||||
def __init__(self, sample: bool = True, chunk_dim: int = 1):
|
||||
super().__init__()
|
||||
self.sample = sample
|
||||
self.chunk_dim = chunk_dim
|
||||
|
||||
def forward(self, z: Tensor) -> Tensor:
|
||||
mean, logvar = torch.chunk(z, 2, dim=self.chunk_dim)
|
||||
# import pdb;pdb.set_trace()
|
||||
if self.sample:
|
||||
std = torch.exp(0.5 * logvar)
|
||||
return mean #+ std * torch.randn_like(mean)
|
||||
else:
|
||||
return mean
|
||||
|
||||
|
||||
class AutoEncoder(nn.Module):
|
||||
def __init__(self, params: AutoEncoderParams):
|
||||
super().__init__()
|
||||
self.encoder = Encoder(
|
||||
resolution=params.resolution,
|
||||
in_channels=params.in_channels,
|
||||
ch=params.ch,
|
||||
ch_mult=params.ch_mult,
|
||||
num_res_blocks=params.num_res_blocks,
|
||||
z_channels=params.z_channels,
|
||||
)
|
||||
self.decoder = Decoder(
|
||||
resolution=params.resolution,
|
||||
in_channels=params.in_channels,
|
||||
ch=params.ch,
|
||||
out_ch=params.out_ch,
|
||||
ch_mult=params.ch_mult,
|
||||
num_res_blocks=params.num_res_blocks,
|
||||
z_channels=params.z_channels,
|
||||
)
|
||||
self.reg = DiagonalGaussian()
|
||||
|
||||
self.scale_factor = params.scale_factor
|
||||
self.shift_factor = params.shift_factor
|
||||
|
||||
def encode(self, x: Tensor) -> Tensor:
|
||||
z = self.reg(self.encoder(x))
|
||||
z = self.scale_factor * (z - self.shift_factor)
|
||||
return z
|
||||
|
||||
def decode(self, z: Tensor) -> Tensor:
|
||||
z = z / self.scale_factor + self.shift_factor
|
||||
return self.decoder(z)
|
||||
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
return self.decode(self.encode(x))
|
||||
@@ -1,47 +0,0 @@
|
||||
from torch import Tensor, nn
|
||||
from transformers import (CLIPTextModel, CLIPTokenizer, T5EncoderModel,
|
||||
T5Tokenizer)
|
||||
|
||||
|
||||
class HFEmbedder(nn.Module):
|
||||
def __init__(self, version: str, max_length: int, is_clip, **hf_kwargs):
|
||||
super().__init__()
|
||||
self.is_clip = is_clip
|
||||
self.max_length = max_length
|
||||
self.output_key = "pooler_output" if self.is_clip else "last_hidden_state"
|
||||
|
||||
if version:
|
||||
#if version == 'black-forest-labs/FLUX.1-dev':
|
||||
if self.is_clip:
|
||||
self.tokenizer: T5Tokenizer = CLIPTokenizer.from_pretrained(version, max_length=max_length, subfolder="tokenizer")
|
||||
self.hf_module: T5EncoderModel = CLIPTextModel.from_pretrained(version,subfolder='text_encoder' , **hf_kwargs)
|
||||
else:
|
||||
self.tokenizer: T5Tokenizer = T5Tokenizer.from_pretrained(version, max_length=max_length, subfolder="tokenizer_2")
|
||||
self.hf_module: T5EncoderModel = T5EncoderModel.from_pretrained(version,subfolder='text_encoder_2' , **hf_kwargs)
|
||||
else:
|
||||
if self.is_clip:
|
||||
self.tokenizer: CLIPTokenizer = CLIPTokenizer.from_pretrained(version, max_length=max_length)
|
||||
self.hf_module: CLIPTextModel = CLIPTextModel.from_pretrained(version, **hf_kwargs)
|
||||
else:
|
||||
self.tokenizer: T5Tokenizer = T5Tokenizer.from_pretrained(version, max_length=max_length)
|
||||
self.hf_module: T5EncoderModel = T5EncoderModel.from_pretrained(version, **hf_kwargs)
|
||||
|
||||
self.hf_module = self.hf_module.eval().requires_grad_(False)
|
||||
|
||||
def forward(self, text: list[str]) -> Tensor:
|
||||
batch_encoding = self.tokenizer(
|
||||
text,
|
||||
truncation=True,
|
||||
max_length=self.max_length,
|
||||
return_length=False,
|
||||
return_overflowing_tokens=False,
|
||||
padding="max_length",
|
||||
return_tensors="pt",
|
||||
)
|
||||
|
||||
outputs = self.hf_module(
|
||||
input_ids=batch_encoding["input_ids"].to(self.hf_module.device),
|
||||
attention_mask=None,
|
||||
output_hidden_states=False,
|
||||
)
|
||||
return outputs[self.output_key]
|
||||
@@ -1,458 +0,0 @@
|
||||
import math
|
||||
from dataclasses import dataclass
|
||||
|
||||
import torch
|
||||
from einops import rearrange
|
||||
from torch import Tensor, nn
|
||||
|
||||
from ..math import attention, rope,apply_rope
|
||||
|
||||
import os
|
||||
|
||||
class EmbedND(nn.Module):
|
||||
def __init__(self, dim: int, theta: int, axes_dim: list[int]):
|
||||
super().__init__()
|
||||
self.dim = dim
|
||||
self.theta = theta
|
||||
self.axes_dim = axes_dim
|
||||
|
||||
def forward(self, ids: Tensor) -> Tensor:
|
||||
n_axes = ids.shape[-1]
|
||||
emb = torch.cat(
|
||||
[rope(ids[..., i], self.axes_dim[i], self.theta) for i in range(n_axes)],
|
||||
dim=-3,
|
||||
)
|
||||
|
||||
return emb.unsqueeze(1)
|
||||
|
||||
|
||||
def timestep_embedding(t: Tensor, dim, max_period=10000, time_factor: float = 1000.0):
|
||||
"""
|
||||
Create sinusoidal timestep embeddings.
|
||||
:param t: a 1-D Tensor of N indices, one per batch element.
|
||||
These may be fractional.
|
||||
:param dim: the dimension of the output.
|
||||
:param max_period: controls the minimum frequency of the embeddings.
|
||||
:return: an (N, D) Tensor of positional embeddings.
|
||||
"""
|
||||
t = time_factor * t
|
||||
half = dim // 2
|
||||
freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to(
|
||||
t.device
|
||||
)
|
||||
|
||||
args = t[:, None].float() * freqs[None]
|
||||
embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
|
||||
if dim % 2:
|
||||
embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
|
||||
if torch.is_floating_point(t):
|
||||
embedding = embedding.to(t)
|
||||
return embedding
|
||||
|
||||
|
||||
class MLPEmbedder(nn.Module):
|
||||
def __init__(self, in_dim: int, hidden_dim: int):
|
||||
super().__init__()
|
||||
self.in_layer = nn.Linear(in_dim, hidden_dim, bias=True)
|
||||
self.silu = nn.SiLU()
|
||||
self.out_layer = nn.Linear(hidden_dim, hidden_dim, bias=True)
|
||||
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
return self.out_layer(self.silu(self.in_layer(x)))
|
||||
|
||||
|
||||
class RMSNorm(torch.nn.Module):
|
||||
def __init__(self, dim: int):
|
||||
super().__init__()
|
||||
self.scale = nn.Parameter(torch.ones(dim))
|
||||
|
||||
def forward(self, x: Tensor):
|
||||
x_dtype = x.dtype
|
||||
x = x.float()
|
||||
rrms = torch.rsqrt(torch.mean(x**2, dim=-1, keepdim=True) + 1e-6)
|
||||
return (x * rrms).to(dtype=x_dtype) * self.scale
|
||||
|
||||
|
||||
class QKNorm(torch.nn.Module):
|
||||
def __init__(self, dim: int):
|
||||
super().__init__()
|
||||
self.query_norm = RMSNorm(dim)
|
||||
self.key_norm = RMSNorm(dim)
|
||||
|
||||
def forward(self, q: Tensor, k: Tensor, v: Tensor) -> tuple[Tensor, Tensor]:
|
||||
q = self.query_norm(q)
|
||||
k = self.key_norm(k)
|
||||
return q.to(v), k.to(v)
|
||||
|
||||
|
||||
class SelfAttention(nn.Module):
|
||||
def __init__(self, dim: int, num_heads: int = 8, qkv_bias: bool = False):
|
||||
super().__init__()
|
||||
self.num_heads = num_heads
|
||||
head_dim = dim // num_heads
|
||||
|
||||
self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
|
||||
self.norm = QKNorm(head_dim)
|
||||
self.proj = nn.Linear(dim, dim)
|
||||
|
||||
def forward(self, x: Tensor, pe: Tensor) -> Tensor:
|
||||
qkv = self.qkv(x)
|
||||
q, k, v = rearrange(qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
|
||||
q, k = self.norm(q, k, v)
|
||||
x = attention(q, k, v, pe=pe)
|
||||
x = self.proj(x)
|
||||
return x
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModulationOut:
|
||||
shift: Tensor
|
||||
scale: Tensor
|
||||
gate: Tensor
|
||||
|
||||
|
||||
class Modulation(nn.Module):
|
||||
def __init__(self, dim: int, double: bool):
|
||||
super().__init__()
|
||||
self.is_double = double
|
||||
self.multiplier = 6 if double else 3
|
||||
self.lin = nn.Linear(dim, self.multiplier * dim, bias=True)
|
||||
|
||||
def forward(self, vec: Tensor) -> tuple[ModulationOut, ModulationOut | None]:
|
||||
out = self.lin(nn.functional.silu(vec))[:, None, :].chunk(self.multiplier, dim=-1)
|
||||
|
||||
return (
|
||||
ModulationOut(*out[:3]),
|
||||
ModulationOut(*out[3:]) if self.is_double else None,
|
||||
)
|
||||
|
||||
|
||||
class DoubleStreamBlock(nn.Module):
|
||||
def __init__(self, hidden_size: int, num_heads: int, mlp_ratio: float, qkv_bias: bool = False):
|
||||
super().__init__()
|
||||
|
||||
mlp_hidden_dim = int(hidden_size * mlp_ratio)
|
||||
self.num_heads = num_heads
|
||||
self.hidden_size = hidden_size
|
||||
self.img_mod = Modulation(hidden_size, double=True)
|
||||
self.img_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
|
||||
self.img_attn = SelfAttention(dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias)
|
||||
|
||||
self.img_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
|
||||
self.img_mlp = nn.Sequential(
|
||||
nn.Linear(hidden_size, mlp_hidden_dim, bias=True),
|
||||
nn.GELU(approximate="tanh"),
|
||||
nn.Linear(mlp_hidden_dim, hidden_size, bias=True),
|
||||
)
|
||||
|
||||
self.txt_mod = Modulation(hidden_size, double=True)
|
||||
self.txt_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
|
||||
self.txt_attn = SelfAttention(dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias)
|
||||
|
||||
self.txt_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
|
||||
self.txt_mlp = nn.Sequential(
|
||||
nn.Linear(hidden_size, mlp_hidden_dim, bias=True),
|
||||
nn.GELU(approximate="tanh"),
|
||||
nn.Linear(mlp_hidden_dim, hidden_size, bias=True),
|
||||
)
|
||||
|
||||
def forward(self, img: Tensor, txt: Tensor, vec: Tensor, pe: Tensor) -> tuple[Tensor, Tensor]:
|
||||
img_mod1, img_mod2 = self.img_mod(vec)
|
||||
txt_mod1, txt_mod2 = self.txt_mod(vec)
|
||||
|
||||
# prepare image for attention
|
||||
img_modulated = self.img_norm1(img)
|
||||
img_modulated = (1 + img_mod1.scale) * img_modulated + img_mod1.shift
|
||||
img_qkv = self.img_attn.qkv(img_modulated)
|
||||
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)
|
||||
|
||||
img_q, img_k = self.img_attn.norm(img_q, img_k, img_v)
|
||||
# prepare txt for attention
|
||||
txt_modulated = self.txt_norm1(txt)
|
||||
txt_modulated = (1 + txt_mod1.scale) * txt_modulated + txt_mod1.shift
|
||||
txt_qkv = self.txt_attn.qkv(txt_modulated)
|
||||
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)
|
||||
txt_q, txt_k = self.txt_attn.norm(txt_q, txt_k, txt_v)
|
||||
|
||||
# run actual attention
|
||||
q = torch.cat((txt_q, img_q), dim=2)
|
||||
k = torch.cat((txt_k, img_k), dim=2)
|
||||
v = torch.cat((txt_v, img_v), dim=2)
|
||||
# import pdb;pdb.set_trace()
|
||||
attn = attention(q, k, v, pe=pe)
|
||||
|
||||
txt_attn, img_attn = attn[:, : txt.shape[1]], attn[:, txt.shape[1] :]
|
||||
|
||||
# calculate the img bloks
|
||||
img = img + img_mod1.gate * self.img_attn.proj(img_attn)
|
||||
img = img + img_mod2.gate * self.img_mlp((1 + img_mod2.scale) * self.img_norm2(img) + img_mod2.shift)
|
||||
|
||||
# calculate the txt bloks
|
||||
txt = txt + txt_mod1.gate * self.txt_attn.proj(txt_attn)
|
||||
txt = txt + txt_mod2.gate * self.txt_mlp((1 + txt_mod2.scale) * self.txt_norm2(txt) + txt_mod2.shift)
|
||||
return img, txt
|
||||
class SingleStreamBlock(nn.Module):
|
||||
"""
|
||||
A DiT block with parallel linear layers as described in
|
||||
https://arxiv.org/abs/2302.05442 and adapted modulation interface.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
hidden_size: int,
|
||||
num_heads: int,
|
||||
mlp_ratio: float = 4.0,
|
||||
qk_scale: float | None = None,
|
||||
):
|
||||
super().__init__()
|
||||
self.hidden_dim = hidden_size
|
||||
self.num_heads = num_heads
|
||||
head_dim = hidden_size // num_heads
|
||||
self.scale = qk_scale or head_dim**-0.5
|
||||
|
||||
self.mlp_hidden_dim = int(hidden_size * mlp_ratio)
|
||||
# qkv and mlp_in
|
||||
self.linear1 = nn.Linear(hidden_size, hidden_size * 3 + self.mlp_hidden_dim)
|
||||
# proj and mlp_out
|
||||
self.linear2 = nn.Linear(hidden_size + self.mlp_hidden_dim, hidden_size)
|
||||
|
||||
self.norm = QKNorm(head_dim)
|
||||
|
||||
self.hidden_size = hidden_size
|
||||
self.pre_norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
|
||||
|
||||
self.mlp_act = nn.GELU(approximate="tanh")
|
||||
self.modulation = Modulation(hidden_size, double=False)
|
||||
|
||||
def forward(self, x: Tensor, vec: Tensor, pe: Tensor) -> Tensor:
|
||||
mod, _ = self.modulation(vec)
|
||||
x_mod = (1 + mod.scale) * self.pre_norm(x) + mod.shift
|
||||
qkv, mlp = torch.split(self.linear1(x_mod), [3 * self.hidden_size, self.mlp_hidden_dim], dim=-1)
|
||||
|
||||
q, k, v = rearrange(qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
|
||||
q, k = self.norm(q, k, v)
|
||||
|
||||
# compute attention
|
||||
attn = attention(q, k, v, pe=pe)
|
||||
# compute activation in mlp stream, cat again and run second linear layer
|
||||
output = self.linear2(torch.cat((attn, self.mlp_act(mlp)), 2))
|
||||
return x + mod.gate * output
|
||||
|
||||
|
||||
class LastLayer(nn.Module):
|
||||
def __init__(self, hidden_size: int, patch_size: int, out_channels: int):
|
||||
super().__init__()
|
||||
self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
|
||||
self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True)
|
||||
self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size, bias=True))
|
||||
|
||||
def forward(self, x: Tensor, vec: Tensor) -> Tensor:
|
||||
shift, scale = self.adaLN_modulation(vec).chunk(2, dim=1)
|
||||
x = (1 + scale[:, None, :]) * self.norm_final(x) + shift[:, None, :]
|
||||
x = self.linear(x)
|
||||
return x
|
||||
|
||||
class LastLayer(nn.Module):
|
||||
def __init__(self, hidden_size: int, patch_size: int, out_channels: int):
|
||||
super().__init__()
|
||||
self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
|
||||
self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True)
|
||||
self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size, bias=True))
|
||||
|
||||
def forward(self, x: Tensor, vec: Tensor) -> Tensor:
|
||||
shift, scale = self.adaLN_modulation(vec).chunk(2, dim=1)
|
||||
x = (1 + scale[:, None, :]) * self.norm_final(x) + shift[:, None, :]
|
||||
x = self.linear(x)
|
||||
return x
|
||||
|
||||
import threading
|
||||
# 线程局部存储，避免多线程冲突
|
||||
_local_storage = threading.local()
|
||||
|
||||
def get_safe_feature_dict():
|
||||
"""获取线程安全的特征字典"""
|
||||
if not hasattr(_local_storage, 'safe_feature_dict'):
|
||||
_local_storage.safe_feature_dict = {}
|
||||
return _local_storage.safe_feature_dict
|
||||
|
||||
def sync_features_to_safe_dict(info):
|
||||
"""将 info['feature'] 中的特征同步到安全字典"""
|
||||
if 'feature' in info:
|
||||
safe_dict = get_safe_feature_dict()
|
||||
for key, value in info['feature'].items():
|
||||
if isinstance(value, torch.Tensor):
|
||||
# 保存张量的副本，避免设备间干扰
|
||||
safe_dict[key] = value.detach() # 保持在当前设备
|
||||
else:
|
||||
safe_dict[key] = value
|
||||
|
||||
def get_feature_from_safe_dict(feature_name, target_device):
|
||||
"""从安全字典中获取特征"""
|
||||
safe_dict = get_safe_feature_dict()
|
||||
if feature_name in safe_dict:
|
||||
feature = safe_dict[feature_name]
|
||||
if isinstance(feature, torch.Tensor):
|
||||
return feature.to(target_device)
|
||||
return feature
|
||||
return None
|
||||
|
||||
def clear_safe_feature_dict():
|
||||
"""清空安全特征字典"""
|
||||
if hasattr(_local_storage, 'safe_feature_dict'):
|
||||
_local_storage.safe_feature_dict.clear()
|
||||
|
||||
|
||||
class DoubleStreamBlock_kv(DoubleStreamBlock):
|
||||
def __init__(self, hidden_size: int, num_heads: int, mlp_ratio: float, qkv_bias: bool = False):
|
||||
super().__init__(hidden_size, num_heads, mlp_ratio, qkv_bias)
|
||||
|
||||
def forward(self, img: Tensor, txt: Tensor, vec: Tensor, pe: Tensor, info) -> tuple[Tensor, Tensor]:
|
||||
|
||||
img_mod1, img_mod2 = self.img_mod(vec)
|
||||
txt_mod1, txt_mod2 = self.txt_mod(vec)
|
||||
|
||||
# prepare image for attention
|
||||
img_modulated = self.img_norm1(img)
|
||||
img_modulated = (1 + img_mod1.scale) * img_modulated + img_mod1.shift
|
||||
img_qkv = self.img_attn.qkv(img_modulated)
|
||||
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)
|
||||
|
||||
img_q, img_k = self.img_attn.norm(img_q, img_k, img_v)
|
||||
# prepare txt for attention
|
||||
txt_modulated = self.txt_norm1(txt)
|
||||
txt_modulated = (1 + txt_mod1.scale) * txt_modulated + txt_mod1.shift
|
||||
txt_qkv = self.txt_attn.qkv(txt_modulated)
|
||||
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)
|
||||
txt_q, txt_k = self.txt_attn.norm(txt_q, txt_k, txt_v)
|
||||
|
||||
feature_k_name = str(info['t']) + '_' + str(info['id']) + '_' + 'MB' + '_' + 'K'
|
||||
feature_v_name = str(info['t']) + '_' + str(info['id']) + '_' + 'MB' + '_' + 'V'
|
||||
#print(info['inverse'])
|
||||
if info['inverse']:
|
||||
info['feature'][feature_k_name] = img_k.cpu()
|
||||
info['feature'][feature_v_name] = img_v.cpu()
|
||||
|
||||
sync_features_to_safe_dict(info)
|
||||
|
||||
q = torch.cat((txt_q, img_q), dim=2)
|
||||
k = torch.cat((txt_k, img_k), dim=2)
|
||||
v = torch.cat((txt_v, img_v), dim=2)
|
||||
if 'attention_mask' in info:
|
||||
if isinstance(info['attention_mask'], torch.Tensor):
|
||||
attention_mask_=info['attention_mask'].to(img.device)
|
||||
else:
|
||||
attention_mask_=info['attention_mask']
|
||||
attn = attention(q, k, v, pe=pe,attention_mask=attention_mask_)
|
||||
else:
|
||||
attn = attention(q, k, v, pe=pe)
|
||||
|
||||
else:
|
||||
if is_feature_available_in_info(info, feature_k_name):
|
||||
source_img_k = info['feature'][feature_k_name].to(img.device)
|
||||
source_img_v = info['feature'][feature_v_name].to(img.device)
|
||||
else:
|
||||
source_img_k = get_feature_from_safe_dict(feature_k_name, img.device)
|
||||
source_img_v = get_feature_from_safe_dict(feature_v_name, img.device)
|
||||
|
||||
mask_indices = info['mask_indices']
|
||||
source_img_k[:, :, mask_indices, ...] = img_k
|
||||
source_img_v[:, :, mask_indices, ...] = img_v
|
||||
|
||||
q = torch.cat((txt_q, img_q), dim=2)
|
||||
k = torch.cat((txt_k, source_img_k), dim=2)
|
||||
v = torch.cat((txt_v, source_img_v), dim=2)
|
||||
#attn = attention(q, k, v, pe=pe, pe_q = info['pe_mask'])
|
||||
attn = attention(q, k, v, pe=pe, pe_q = info['pe_mask'],attention_mask=info['attention_scale'])
|
||||
|
||||
txt_attn, img_attn = attn[:, : txt.shape[1]], attn[:, txt.shape[1] :]
|
||||
|
||||
# calculate the img bloks
|
||||
img = img + img_mod1.gate * self.img_attn.proj(img_attn)
|
||||
img = img + img_mod2.gate * self.img_mlp((1 + img_mod2.scale) * self.img_norm2(img) + img_mod2.shift)
|
||||
|
||||
# calculate the txt bloks
|
||||
txt = txt + txt_mod1.gate * self.txt_attn.proj(txt_attn)
|
||||
txt = txt + txt_mod2.gate * self.txt_mlp((1 + txt_mod2.scale) * self.txt_norm2(txt) + txt_mod2.shift)
|
||||
return img, txt
|
||||
|
||||
def is_feature_available_in_info(info: dict, feature_name: str) -> bool:
|
||||
"""检查特征是否在 info 字典中可用"""
|
||||
try:
|
||||
return (feature_name in info.get('feature', {}) and
|
||||
info['feature'][feature_name] is not None)
|
||||
except:
|
||||
return False
|
||||
|
||||
class SingleStreamBlock_kv(SingleStreamBlock):
|
||||
"""
|
||||
A DiT block with parallel linear layers as described in
|
||||
https://arxiv.org/abs/2302.05442 and adapted modulation interface.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
hidden_size: int,
|
||||
num_heads: int,
|
||||
mlp_ratio: float = 4.0,
|
||||
qk_scale: float | None = None,
|
||||
):
|
||||
super().__init__(hidden_size, num_heads, mlp_ratio, qk_scale)
|
||||
|
||||
def forward(self,x: Tensor, vec: Tensor, pe: Tensor, info,) -> Tensor:
|
||||
token_len=info['token_len']
|
||||
mod, _ = self.modulation(vec)
|
||||
x_mod = (1 + mod.scale) * self.pre_norm(x) + mod.shift
|
||||
qkv, mlp = torch.split(self.linear1(x_mod), [3 * self.hidden_size, self.mlp_hidden_dim], dim=-1)
|
||||
|
||||
q, k, v = rearrange(qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
|
||||
q, k = self.norm(q, k, v)
|
||||
# img_k = k[:, :, 512:, ...]
|
||||
# img_v = v[:, :, 512:, ...]
|
||||
|
||||
# txt_k = k[:, :, :512, ...]
|
||||
# txt_v = v[:, :, :512, ...]
|
||||
img_k = k[:, :, token_len:, ...] #令牌数与clip的必须一致，否则会报错，加载comfyUI的量化T5
|
||||
img_v = v[:, :, token_len:, ...]
|
||||
|
||||
txt_k = k[:, :, :token_len, ...]
|
||||
txt_v = v[:, :, :token_len, ...]
|
||||
|
||||
|
||||
feature_k_name = str(info['t']) + '_' + str(info['id']) + '_' + 'SB' + '_' + 'K'
|
||||
feature_v_name = str(info['t']) + '_' + str(info['id']) + '_' + 'SB' + '_' + 'V'
|
||||
if info['inverse']:
|
||||
info['feature'][feature_k_name] = img_k.cpu()
|
||||
info['feature'][feature_v_name] = img_v.cpu()
|
||||
|
||||
sync_features_to_safe_dict(info)
|
||||
|
||||
if 'attention_mask' in info:
|
||||
if isinstance(info['attention_mask'], torch.Tensor):
|
||||
attention_mask_=info['attention_mask'].to(x.device)
|
||||
else:
|
||||
attention_mask_=info['attention_mask']
|
||||
attn = attention(q, k, v, pe=pe,attention_mask=attention_mask_)
|
||||
else:
|
||||
attn = attention(q, k, v, pe=pe)
|
||||
|
||||
else:
|
||||
if is_feature_available_in_info(info, feature_k_name):
|
||||
source_img_k = info['feature'][feature_k_name].to(x.device)
|
||||
source_img_v = info['feature'][feature_v_name].to(x.device)
|
||||
else:
|
||||
source_img_k = get_feature_from_safe_dict(feature_k_name, x.device)
|
||||
source_img_v = get_feature_from_safe_dict(feature_v_name, x.device)
|
||||
|
||||
mask_indices = info['mask_indices']
|
||||
#print(source_img_k.shape,img_k.shape,mask_indices.shape) #torch.Size([1, 24, 768, 128]) torch.Size([1, 24, 0, 128]) torch.Size([152])
|
||||
source_img_k[:, :, mask_indices, ...] = img_k # [1, 24, 154, 128] [24, 0, 128]
|
||||
source_img_v[:, :, mask_indices, ...] = img_v
|
||||
|
||||
k = torch.cat((txt_k, source_img_k), dim=2)
|
||||
v = torch.cat((txt_v, source_img_v), dim=2)
|
||||
#attn = attention(q, k, v, pe=pe, pe_q = info['pe_mask'])
|
||||
attn = attention(q, k, v, pe=pe, pe_q = info['pe_mask'],attention_mask=info['attention_scale'])
|
||||
|
||||
# compute activation in mlp stream, cat again and run second linear layer
|
||||
output = self.linear2(torch.cat((attn, self.mlp_act(mlp)), 2))
|
||||
return x + mod.gate * output
|
||||
@@ -1,292 +0,0 @@
|
||||
import math
|
||||
from typing import Callable
|
||||
|
||||
import torch
|
||||
from einops import rearrange, repeat
|
||||
from torch import Tensor
|
||||
from .modules.layers import clear_safe_feature_dict
|
||||
from .model import Flux,Flux_kv
|
||||
from .modules.conditioner import HFEmbedder
|
||||
from tqdm import tqdm
|
||||
from tqdm.contrib import tzip
|
||||
|
||||
def get_noise(
|
||||
num_samples: int,
|
||||
height: int,
|
||||
width: int,
|
||||
device: torch.device,
|
||||
dtype: torch.dtype,
|
||||
seed: int,
|
||||
):
|
||||
return torch.randn(
|
||||
num_samples,
|
||||
16,
|
||||
# allow for packing
|
||||
2 * math.ceil(height / 16),
|
||||
2 * math.ceil(width / 16),
|
||||
device=device,
|
||||
dtype=dtype,
|
||||
generator=torch.Generator(device=device).manual_seed(seed),
|
||||
)
|
||||
|
||||
def prepare(clip,device,img: Tensor, prompt: str | list[str]) -> dict[str, Tensor]:
|
||||
bs, c, h, w = img.shape
|
||||
if bs == 1 and not isinstance(prompt, str):
|
||||
bs = len(prompt)
|
||||
|
||||
img = rearrange(img, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2)
|
||||
if img.shape[0] == 1 and bs > 1:
|
||||
img = repeat(img, "1 ... -> bs ...", bs=bs)
|
||||
|
||||
img_ids = torch.zeros(h // 2, w // 2, 3)
|
||||
img_ids[..., 1] = img_ids[..., 1] + torch.arange(h // 2)[:, None]
|
||||
img_ids[..., 2] = img_ids[..., 2] + torch.arange(w // 2)[None, :]
|
||||
img_ids = repeat(img_ids, "h w c -> b (h w) c", b=bs)
|
||||
|
||||
# if isinstance(prompt, str):
|
||||
# prompt = [prompt]
|
||||
|
||||
# txt = t5(prompt)
|
||||
clip.tokenizer.min_length=512
|
||||
tokens = clip.tokenize(prompt)
|
||||
tokens["t5xxl"] = clip.tokenize(prompt)["t5xxl"]
|
||||
txt = clip.encode_from_tokens(tokens, return_pooled=True, return_dict=True).pop("cond")
|
||||
|
||||
|
||||
if txt.shape[0] == 1 and bs > 1:
|
||||
txt = repeat(txt, "1 ... -> bs ...", bs=bs)
|
||||
txt_ids = torch.zeros(bs, txt.shape[1], 3)
|
||||
|
||||
#vec = clip(prompt)
|
||||
tokens["l"] = clip.tokenize(prompt)["l"]
|
||||
vec = clip.encode_from_tokens(tokens, return_dict=True).pop("pooled_output")
|
||||
|
||||
#print(vec.shape, img.shape, txt.shape, img_ids.shape, txt_ids.shape) # torch.Size([1, 768]) torch.Size([1, 1024, 64]) torch.Size([1, 512, 4096]) torch.Size([1, 1024, 3]) torch.Size([1, 512, 3])
|
||||
if vec.shape[0] == 1 and bs > 1:
|
||||
vec = repeat(vec, "1 ... -> bs ...", bs=bs)
|
||||
|
||||
return {
|
||||
"img": img.to(device,torch.bfloat16),
|
||||
"img_ids": img_ids.to(device,torch.bfloat16),
|
||||
"txt": txt.to(device,torch.bfloat16),
|
||||
"txt_ids": txt_ids.to(device,torch.bfloat16),
|
||||
"vec": vec.to(device,torch.bfloat16),
|
||||
}
|
||||
|
||||
def prepare_(t5: HFEmbedder, clip: HFEmbedder,device, img: Tensor, prompt: str | list[str]) -> dict[str, Tensor]:
|
||||
bs, c, h, w = img.shape
|
||||
if bs == 1 and not isinstance(prompt, str):
|
||||
bs = len(prompt)
|
||||
|
||||
img = rearrange(img, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2)
|
||||
if img.shape[0] == 1 and bs > 1:
|
||||
img = repeat(img, "1 ... -> bs ...", bs=bs)
|
||||
|
||||
img_ids = torch.zeros(h // 2, w // 2, 3)
|
||||
img_ids[..., 1] = img_ids[..., 1] + torch.arange(h // 2)[:, None]
|
||||
img_ids[..., 2] = img_ids[..., 2] + torch.arange(w // 2)[None, :]
|
||||
img_ids = repeat(img_ids, "h w c -> b (h w) c", b=bs)
|
||||
|
||||
if isinstance(prompt, str):
|
||||
prompt = [prompt]
|
||||
txt = t5(prompt)
|
||||
if txt.shape[0] == 1 and bs > 1:
|
||||
txt = repeat(txt, "1 ... -> bs ...", bs=bs)
|
||||
txt_ids = torch.zeros(bs, txt.shape[1], 3)
|
||||
|
||||
vec = clip(prompt)
|
||||
if vec.shape[0] == 1 and bs > 1:
|
||||
vec = repeat(vec, "1 ... -> bs ...", bs=bs)
|
||||
#print(img.shape,vec.shape, txt.shape, img_ids.shape) #torch.Size([1, 1024, 64]) torch.Size([1, 768]) torch.Size([1, 512, 4096]) torch.Size([1, 1024, 3])
|
||||
return {
|
||||
"img": img.to(device,torch.bfloat16),
|
||||
"img_ids": img_ids.to(device,torch.bfloat16),
|
||||
"txt": txt.to(device,torch.bfloat16),
|
||||
"txt_ids": txt_ids.to(device,torch.bfloat16),
|
||||
"vec": vec.to(device,torch.bfloat16),
|
||||
}
|
||||
|
||||
def time_shift(mu: float, sigma: float, t: Tensor):
|
||||
return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma)
|
||||
|
||||
|
||||
def get_lin_function(
|
||||
x1: float = 256, y1: float = 0.5, x2: float = 4096, y2: float = 1.15
|
||||
) -> Callable[[float], float]:
|
||||
m = (y2 - y1) / (x2 - x1)
|
||||
b = y1 - m * x1
|
||||
return lambda x: m * x + b
|
||||
|
||||
|
||||
def get_schedule(
|
||||
num_steps: int,
|
||||
image_seq_len: int,
|
||||
base_shift: float = 0.5,
|
||||
max_shift: float = 1.15,
|
||||
shift: bool = True,
|
||||
) -> list[float]:
|
||||
# extra step for zero
|
||||
timesteps = torch.linspace(1, 0, num_steps + 1)
|
||||
|
||||
# shifting the schedule to favor high timesteps for higher signal images
|
||||
if shift:
|
||||
# estimate mu based on linear estimation between two points
|
||||
mu = get_lin_function(y1=base_shift, y2=max_shift)(image_seq_len)
|
||||
timesteps = time_shift(mu, 1.0, timesteps)
|
||||
|
||||
return timesteps.tolist()
|
||||
|
||||
|
||||
def denoise(
|
||||
model: Flux,
|
||||
# model input
|
||||
img: Tensor,
|
||||
img_ids: Tensor,
|
||||
txt: Tensor,
|
||||
txt_ids: Tensor,
|
||||
vec: Tensor,
|
||||
# sampling parameters
|
||||
timesteps: list[float],
|
||||
guidance: float = 4.0,
|
||||
):
|
||||
# this is ignored for schnell
|
||||
guidance_vec = torch.full((img.shape[0],), guidance, device=img.device, dtype=img.dtype)
|
||||
for i, (t_curr, t_prev) in enumerate(zip(timesteps[:-1], timesteps[1:])):
|
||||
t_vec = torch.full((img.shape[0],), t_curr, dtype=img.dtype, device=img.device)
|
||||
pred = model(
|
||||
img=img,
|
||||
img_ids=img_ids,
|
||||
txt=txt,
|
||||
txt_ids=txt_ids,
|
||||
y=vec,
|
||||
timesteps=t_vec,
|
||||
guidance=guidance_vec,
|
||||
)
|
||||
|
||||
img = img + (t_prev - t_curr) * pred
|
||||
|
||||
return img
|
||||
|
||||
def unpack(x: Tensor, height: int, width: int) -> Tensor:
|
||||
return rearrange(
|
||||
x,
|
||||
"b (h w) (c ph pw) -> b c (h ph) (w pw)",
|
||||
h=math.ceil(height / 16),
|
||||
w=math.ceil(width / 16),
|
||||
ph=2,
|
||||
pw=2,
|
||||
)
|
||||
|
||||
def denoise_kv(
|
||||
model: Flux_kv,
|
||||
# model input
|
||||
img: Tensor,
|
||||
img_ids: Tensor,
|
||||
txt: Tensor,
|
||||
txt_ids: Tensor,
|
||||
vec: Tensor,
|
||||
# sampling parameters
|
||||
timesteps: list[float],
|
||||
inverse,
|
||||
info,
|
||||
guidance: float = 4.0
|
||||
):
|
||||
|
||||
if inverse:
|
||||
timesteps = timesteps[::-1]
|
||||
|
||||
guidance_vec = torch.full((img.shape[0],), guidance, device=img.device, dtype=img.dtype)
|
||||
|
||||
for i, (t_curr, t_prev) in enumerate(tzip(timesteps[:-1], timesteps[1:])):
|
||||
t_vec = torch.full((img.shape[0],), t_curr, dtype=img.dtype, device=img.device)
|
||||
info['t'] = t_prev if inverse else t_curr
|
||||
|
||||
if inverse:
|
||||
img_name = str(info['t']) + '_' + 'img'
|
||||
info['feature'][img_name] = img.cpu()
|
||||
else:
|
||||
img_name = str(info['t']) + '_' + 'img'
|
||||
source_img = info['feature'][img_name].to(img.device)
|
||||
img = source_img[:, info['mask_indices'],...] * (1 - info['mask'][:, info['mask_indices'],...]) + img * info['mask'][:, info['mask_indices'],...]
|
||||
pred = model(
|
||||
img=img,
|
||||
img_ids=img_ids,
|
||||
txt=txt,
|
||||
txt_ids=txt_ids,
|
||||
y=vec,
|
||||
timesteps=t_vec,
|
||||
guidance=guidance_vec,
|
||||
info=info
|
||||
)
|
||||
img = img + (t_prev - t_curr) * pred
|
||||
return img, info
|
||||
|
||||
def denoise_kv_inf(
|
||||
model: Flux_kv,
|
||||
# model input
|
||||
img: Tensor,
|
||||
img_ids: Tensor,
|
||||
source_txt: Tensor,
|
||||
source_txt_ids: Tensor,
|
||||
source_vec: Tensor,
|
||||
target_txt: Tensor,
|
||||
target_txt_ids: Tensor,
|
||||
target_vec: Tensor,
|
||||
# sampling parameters
|
||||
timesteps: list[float],
|
||||
target_guidance: float = 4.0,
|
||||
source_guidance: float = 4.0,
|
||||
info: dict = {},
|
||||
):
|
||||
|
||||
target_guidance_vec = torch.full((img.shape[0],), target_guidance, device=img.device, dtype=img.dtype)
|
||||
source_guidance_vec = torch.full((img.shape[0],), source_guidance, device=img.device, dtype=img.dtype)
|
||||
|
||||
mask_indices = info['mask_indices']
|
||||
init_img = img.clone()
|
||||
z_fe = img[:, mask_indices,...]
|
||||
|
||||
noise_list = []
|
||||
for i in range(len(timesteps)):
|
||||
noise = torch.randn(init_img.size(), dtype=init_img.dtype,
|
||||
layout=init_img.layout, device=init_img.device,
|
||||
generator=torch.Generator(device=init_img.device).manual_seed(0))
|
||||
noise_list.append(noise)
|
||||
|
||||
for i, (t_curr, t_prev) in enumerate(tzip(timesteps[:-1], timesteps[1:])):
|
||||
|
||||
info['t'] = t_curr
|
||||
t_vec = torch.full((img.shape[0],), t_curr, dtype=img.dtype, device=img.device)
|
||||
|
||||
z_src = (1 - t_curr) * init_img + t_curr * noise_list[i]
|
||||
z_tar = z_src[:, mask_indices,...] - init_img[:, mask_indices,...] + z_fe
|
||||
|
||||
info['inverse'] = True
|
||||
info['feature'] = {}
|
||||
v_src = model(
|
||||
img=z_src,
|
||||
img_ids=img_ids,
|
||||
txt=source_txt,
|
||||
txt_ids=source_txt_ids,
|
||||
y=source_vec,
|
||||
timesteps=t_vec,
|
||||
guidance=source_guidance_vec,
|
||||
info=info
|
||||
)
|
||||
|
||||
info['inverse'] = False
|
||||
v_tar = model(
|
||||
img=z_tar,
|
||||
img_ids=img_ids,
|
||||
txt=target_txt,
|
||||
txt_ids=target_txt_ids,
|
||||
y=target_vec,
|
||||
timesteps=t_vec,
|
||||
guidance=target_guidance_vec,
|
||||
info=info
|
||||
)
|
||||
|
||||
v_fe = v_tar - v_src[:, mask_indices,...]
|
||||
z_fe = z_fe + (t_prev - t_curr) * v_fe * info['mask'][:, mask_indices,...]
|
||||
clear_safe_feature_dict()
|
||||
return z_fe, info
|
||||
-333
@@ -1,333 +0,0 @@
|
||||
import os
|
||||
from dataclasses import dataclass
|
||||
import gc
|
||||
from shutil import copy
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from einops import rearrange
|
||||
from huggingface_hub import hf_hub_download
|
||||
#from imwatermark import WatermarkEncoder
|
||||
from safetensors.torch import load_file as load_sft
|
||||
from .modules.layers import MLPEmbedder,DoubleStreamBlock_kv,DoubleStreamBlock,SingleStreamBlock,SingleStreamBlock_kv
|
||||
from .model import Flux, FluxParams
|
||||
from .modules.autoencoder import AutoEncoder, AutoEncoderParams
|
||||
from .modules.conditioner import HFEmbedder
|
||||
import comfy.model_management
|
||||
import json
|
||||
import types # 新增导入
|
||||
import torch.nn.functional as F # 确保F.interpolate可用
|
||||
@dataclass
|
||||
class ModelSpec:
|
||||
params: FluxParams
|
||||
ae_params: AutoEncoderParams
|
||||
ckpt_path: str | None
|
||||
ae_path: str | None
|
||||
repo_id: str | None
|
||||
repo_flow: str | None
|
||||
repo_ae: str | None
|
||||
|
||||
configs = {
|
||||
"flux-dev": ModelSpec(
|
||||
repo_id="black-forest-labs/FLUX.1-dev",
|
||||
repo_flow="flux1-dev.safetensors",
|
||||
repo_ae="ae.safetensors",
|
||||
ckpt_path=os.getenv("FLUX_DEV"),
|
||||
params=FluxParams(
|
||||
in_channels=64,
|
||||
vec_in_dim=768,
|
||||
context_in_dim=4096,
|
||||
hidden_size=3072,
|
||||
mlp_ratio=4.0,
|
||||
num_heads=24,
|
||||
depth=19,
|
||||
depth_single_blocks=38,
|
||||
axes_dim=[16, 56, 56],
|
||||
theta=10_000,
|
||||
qkv_bias=True,
|
||||
guidance_embed=True,
|
||||
),
|
||||
ae_path=os.getenv("AE"),
|
||||
ae_params=AutoEncoderParams(
|
||||
resolution=256,
|
||||
in_channels=3,
|
||||
ch=128,
|
||||
out_ch=3,
|
||||
ch_mult=[1, 2, 4, 4],
|
||||
num_res_blocks=2,
|
||||
z_channels=16,
|
||||
scale_factor=0.3611,
|
||||
shift_factor=0.1159,
|
||||
),
|
||||
),
|
||||
"flux-schnell": ModelSpec(
|
||||
repo_id="black-forest-labs/FLUX.1-schnell",
|
||||
repo_flow="flux1-schnell.safetensors",
|
||||
repo_ae="ae.safetensors",
|
||||
ckpt_path=os.getenv("FLUX_SCHNELL"),
|
||||
params=FluxParams(
|
||||
in_channels=64,
|
||||
vec_in_dim=768,
|
||||
context_in_dim=4096,
|
||||
hidden_size=3072,
|
||||
mlp_ratio=4.0,
|
||||
num_heads=24,
|
||||
depth=19,
|
||||
depth_single_blocks=38,
|
||||
axes_dim=[16, 56, 56],
|
||||
theta=10_000,
|
||||
qkv_bias=True,
|
||||
guidance_embed=False,
|
||||
),
|
||||
ae_path=os.getenv("AE"),
|
||||
ae_params=AutoEncoderParams(
|
||||
resolution=256,
|
||||
in_channels=3,
|
||||
ch=128,
|
||||
out_ch=3,
|
||||
ch_mult=[1, 2, 4, 4],
|
||||
num_res_blocks=2,
|
||||
z_channels=16,
|
||||
scale_factor=0.3611,
|
||||
shift_factor=0.1159,
|
||||
),
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
def print_load_warning(missing: list[str], unexpected: list[str]) -> None:
|
||||
if len(missing) > 0 and len(unexpected) > 0:
|
||||
print(f"Got {len(missing)} missing keys:\n\t" + "\n\t".join(missing))
|
||||
print("\n" + "-" * 79 + "\n")
|
||||
print(f"Got {len(unexpected)} unexpected keys:\n\t" + "\n\t".join(unexpected))
|
||||
elif len(missing) > 0:
|
||||
print(f"Got {len(missing)} missing keys:\n\t" + "\n\t".join(missing))
|
||||
elif len(unexpected) > 0:
|
||||
print(f"Got {len(unexpected)} unexpected keys:\n\t" + "\n\t".join(unexpected))
|
||||
|
||||
|
||||
# 修改加载函数
|
||||
def load_flux_model(ckpt_path, device, flux_cls=Flux):
|
||||
|
||||
# with torch.device("meta" if ckpt_path is not None else device):
|
||||
# model = flux_cls(configs["flux-dev"].params).to(torch.bfloat16)
|
||||
|
||||
from contextlib import nullcontext
|
||||
try:
|
||||
from accelerate import init_empty_weights,load_checkpoint_and_dispatch
|
||||
is_accelerate_available = True
|
||||
except:
|
||||
is_accelerate_available = False
|
||||
ctx = init_empty_weights if is_accelerate_available else nullcontext
|
||||
with ctx():
|
||||
model = flux_cls(configs["flux-dev"].params).to(torch.bfloat16)
|
||||
|
||||
from comfy.model_management import total_vram,total_ram
|
||||
print("Total VRAM {:0.0f} MB, total RAM {:0.0f} MB".format(total_vram, total_ram))
|
||||
max_memory=int(total_vram/1000.0)-2 # 设置最大显存
|
||||
|
||||
model = load_checkpoint_and_dispatch(
|
||||
model,
|
||||
ckpt_path,
|
||||
device_map="auto", # 自动分配设备
|
||||
max_memory={0: f"{max_memory}GiB", "cpu": "20GiB"}, # 指定每个设备的最大内存
|
||||
offload_folder="offload", # 磁盘卸载文件夹
|
||||
no_split_module_classes=["DoubleStreamBlock_kv", "SingleStreamBlock_kv"],
|
||||
dtype=torch.bfloat16
|
||||
)
|
||||
# if ckpt_path is not None:
|
||||
# print("Loading checkpoint")
|
||||
# # load_sft doesn't support torch.device
|
||||
# sd = load_sft(ckpt_path, device=str(device))
|
||||
# print("Loaded checkpoint",sd.keys())
|
||||
# if "fp8" in ckpt_path:
|
||||
# if is_accelerate_available:
|
||||
# # 使用 set_module_tensor_to_device 逐个设置参数
|
||||
# for name, param in sd.items():
|
||||
# set_module_tensor_to_device(model, name, "cuda", value=param)
|
||||
# else:
|
||||
# # 否则使用常规方式
|
||||
# model.load_state_dict(sd)
|
||||
# # from optimum.quanto import requantize
|
||||
# # import folder_paths
|
||||
# # json_path = os.path.join(folder_paths.base_path, "custom_nodes/ComfyUI_KV_Edit/flux/config.json") #config is for pass block
|
||||
# # with open(json_path,'r') as f:
|
||||
# # quantization_map = json.load(f)
|
||||
# # print(f"Start fp8 requantization process...")
|
||||
# # requantize(model, sd, quantization_map, device=device)
|
||||
# # print("Model is requantized!")
|
||||
# model.load_state_dict(sd, strict=False,)
|
||||
# else:
|
||||
# missing, unexpected = model.load_state_dict(sd, strict=False, assign=True)
|
||||
# print_load_warning(missing, unexpected)
|
||||
# del sd
|
||||
# torch.cuda.empty_cache()
|
||||
# #print_load_warning(missing, unexpected)
|
||||
|
||||
return model
|
||||
|
||||
|
||||
# def load_flux_model_cf(cf_model,device, flux_cls=Flux):
|
||||
# original = cf_model.model.diffusion_model
|
||||
|
||||
# # 保持原参数引用（不复制）
|
||||
# params = original.params
|
||||
# hidden_size = original.hidden_size
|
||||
# num_heads = original.num_heads
|
||||
|
||||
# # 动态创建新模块（保持量化参数）
|
||||
# new_blocks = nn.ModuleList([
|
||||
# DoubleStreamBlock_kv(
|
||||
# hidden_size,
|
||||
# num_heads,
|
||||
# mlp_ratio=params.mlp_ratio,
|
||||
# qkv_bias=params.qkv_bias
|
||||
# ) for _ in range(params.depth)
|
||||
# ])
|
||||
|
||||
# # 参数嫁接（仅替换block参数）
|
||||
# for new_block, old_block in zip(new_blocks, original.double_blocks):
|
||||
# new_block.load_state_dict(old_block.state_dict(), assign=True)
|
||||
|
||||
# # 原子替换（保持原模型内存布局）
|
||||
# original.__class__ = flux_cls
|
||||
# original.double_blocks = new_blocks
|
||||
# original.single_blocks = nn.ModuleList([
|
||||
# SingleStreamBlock_kv(hidden_size, num_heads, params.mlp_ratio)
|
||||
# for _ in range(params.depth_single_blocks)
|
||||
# ])
|
||||
|
||||
# return original
|
||||
|
||||
def load_flux_model_(cf_model, device, flux_cls=Flux):
|
||||
original_sd = cf_model.model.diffusion_model.state_dict()
|
||||
print("Loaded checkpoint",original_sd.keys())
|
||||
del cf_model
|
||||
gc.collect()
|
||||
new_model = flux_cls(configs["flux-dev"].params).to(torch.bfloat16)
|
||||
new_model.load_state_dict(original_sd, strict=False)
|
||||
del original_sd
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
comfy.model_management.unload_all_models()
|
||||
comfy.model_management.soft_empty_cache()
|
||||
#del cf_model.model.diffusion_model
|
||||
torch.cuda.empty_cache()
|
||||
return new_model
|
||||
|
||||
def load_flow_model(name: str, device: str | torch.device = "cuda", hf_download: bool = True, flux_cls=Flux) -> Flux:
|
||||
# Loading Flux
|
||||
print("Init model")
|
||||
|
||||
ckpt_path = configs[name].ckpt_path
|
||||
if (
|
||||
ckpt_path is None
|
||||
and configs[name].repo_id is not None
|
||||
and configs[name].repo_flow is not None
|
||||
and hf_download
|
||||
):
|
||||
ckpt_path = hf_hub_download(configs[name].repo_id, configs[name].repo_flow)
|
||||
|
||||
with torch.device("meta" if ckpt_path is not None else device):
|
||||
model = flux_cls(configs[name].params).to(torch.bfloat16)
|
||||
|
||||
if ckpt_path is not None:
|
||||
print("Loading checkpoint")
|
||||
# load_sft doesn't support torch.device
|
||||
sd = load_sft(ckpt_path, device=str(device))
|
||||
missing, unexpected = model.load_state_dict(sd, strict=False, assign=True)
|
||||
del sd
|
||||
torch.cuda.empty_cache()
|
||||
#print_load_warning(missing, unexpected)
|
||||
return model
|
||||
|
||||
|
||||
def load_t5_(repo,device: str | torch.device = "cuda", max_length: int = 512) -> HFEmbedder:
|
||||
# max length 64, 128, 256 and 512 should work (if your sequence is short enough)
|
||||
# return HFEmbedder("black-forest-labs/FLUX.1-dev", max_length=max_length, is_clip=False, torch_dtype=torch.bfloat16).to(device)
|
||||
return HFEmbedder(repo, max_length=max_length, is_clip=False, torch_dtype=torch.bfloat16).to(device)
|
||||
|
||||
|
||||
def load_clip_(repo,device: str | torch.device = "cuda") -> HFEmbedder:
|
||||
# return HFEmbedder("black-forest-labs/FLUX.1-dev", max_length=77, is_clip=True, torch_dtype=torch.bfloat16).to(device)
|
||||
return HFEmbedder('F:/test/ComfyUI/models/diffusers/black-forest-labs/FLUX.1-dev', max_length=77, is_clip=True, torch_dtype=torch.bfloat16).to(device)
|
||||
|
||||
def load_t5(device: str | torch.device = "cuda", max_length: int = 512) -> HFEmbedder:
|
||||
# max length 64, 128, 256 and 512 should work (if your sequence is short enough)
|
||||
# return HFEmbedder("black-forest-labs/FLUX.1-dev", max_length=max_length, is_clip=False, torch_dtype=torch.bfloat16).to(device)
|
||||
return HFEmbedder("google/t5-v1_1-xxl", max_length=max_length, is_clip=False, torch_dtype=torch.bfloat16).to(device)
|
||||
|
||||
|
||||
def load_clip(device: str | torch.device = "cuda") -> HFEmbedder:
|
||||
# return HFEmbedder("black-forest-labs/FLUX.1-dev", max_length=77, is_clip=True, torch_dtype=torch.bfloat16).to(device)
|
||||
return HFEmbedder("openai/clip-vit-large-patch14", max_length=77, is_clip=True, torch_dtype=torch.bfloat16).to(device)
|
||||
|
||||
|
||||
def load_ae(ckpt_path: str, device: str | torch.device = "cuda", hf_download: bool = True) -> AutoEncoder:
|
||||
# ckpt_path = configs[name].ae_path
|
||||
# if (
|
||||
# ckpt_path is None
|
||||
# and configs[name].repo_id is not None
|
||||
# and configs[name].repo_ae is not None
|
||||
# and hf_download
|
||||
# ):
|
||||
# ckpt_path = hf_hub_download(configs[name].repo_id, configs[name].repo_ae)
|
||||
|
||||
# Loading the autoencoder
|
||||
print("Init AE")
|
||||
#with torch.device("meta" if ckpt_path is not None else device):
|
||||
ae = AutoEncoder(configs["flux-dev"].ae_params).to(torch.bfloat16)
|
||||
|
||||
if ckpt_path is not None:
|
||||
sd = load_sft(ckpt_path, device="cpu")
|
||||
missing, unexpected = ae.load_state_dict(sd, strict=False, assign=True)
|
||||
#print_load_warning(missing, unexpected)
|
||||
del sd
|
||||
torch.cuda.empty_cache()
|
||||
return ae
|
||||
|
||||
def load_ae_cf(ckpt, device: str | torch.device = "cuda", hf_download: bool = True) -> AutoEncoder:
|
||||
# ckpt_path = configs[name].ae_path
|
||||
# if (
|
||||
# ckpt_path is None
|
||||
# and configs[name].repo_id is not None
|
||||
# and configs[name].repo_ae is not None
|
||||
# and hf_download
|
||||
# ):
|
||||
# ckpt_path = hf_hub_download(configs[name].repo_id, configs[name].repo_ae)
|
||||
|
||||
# Loading the autoencoder
|
||||
print("Init AE")
|
||||
#with torch.device("meta" if ckpt_path is not None else device):
|
||||
ae = AutoEncoder(configs["flux-dev"].ae_params).to(torch.bfloat16)
|
||||
|
||||
if ckpt is not None:
|
||||
sd = ckpt.get_sd()
|
||||
missing, unexpected = ae.load_state_dict(sd, strict=False, assign=True)
|
||||
#print_load_warning(missing, unexpected)
|
||||
del sd
|
||||
torch.cuda.empty_cache()
|
||||
return ae
|
||||
|
||||
def load_ae_(name: str, device: str | torch.device = "cuda", hf_download: bool = True) -> AutoEncoder:
|
||||
ckpt_path = configs[name].ae_path
|
||||
if (
|
||||
ckpt_path is None
|
||||
and configs[name].repo_id is not None
|
||||
and configs[name].repo_ae is not None
|
||||
and hf_download
|
||||
):
|
||||
ckpt_path = hf_hub_download(configs[name].repo_id, configs[name].repo_ae)
|
||||
|
||||
# Loading the autoencoder
|
||||
print("Init AE")
|
||||
with torch.device("meta" if ckpt_path is not None else device):
|
||||
ae = AutoEncoder(configs[name].ae_params)
|
||||
|
||||
if ckpt_path is not None:
|
||||
sd = load_sft(ckpt_path, device=str(device))
|
||||
missing, unexpected = ae.load_state_dict(sd, strict=False, assign=True)
|
||||
#print_load_warning(missing, unexpected)
|
||||
del sd
|
||||
torch.cuda.empty_cache()
|
||||
return ae
|
||||
|
||||
@@ -1,106 +0,0 @@
|
||||
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
import torch
|
||||
|
||||
@dataclass
|
||||
class SamplingOptions:
|
||||
source_prompt: str = ''
|
||||
target_prompt: str = ''
|
||||
width: int = 1366
|
||||
height: int = 768
|
||||
inversion_num_steps: int = 0
|
||||
denoise_num_steps: int = 0
|
||||
skip_step: int = 0
|
||||
inversion_guidance: float = 1.0
|
||||
denoise_guidance: float = 1.0
|
||||
seed: int = 42
|
||||
re_init: bool = False
|
||||
attn_mask: bool = False
|
||||
attn_scale: float = 1.0
|
||||
|
||||
class FluxEditor_kv_Wrapper:
|
||||
def __init__(self,offload,device):
|
||||
|
||||
self.device = device
|
||||
self.offload = offload
|
||||
#device_val = "cpu" if self.offload else self.device
|
||||
# self.model=Flux_kv_edit(device_val,model_path,name='flux-dev')
|
||||
# self.model.eval()
|
||||
self.info = {}
|
||||
# if self.offload:
|
||||
# self.model.cpu()
|
||||
# torch.cuda.empty_cache()
|
||||
# self.ae.encoder.to(self.device)
|
||||
|
||||
@torch.inference_mode()
|
||||
def inverse(self,model,opts,mask,inp):
|
||||
self.z0 = None
|
||||
self.zt = None
|
||||
|
||||
if 'feature' in self.info:
|
||||
key_list = list(self.info['feature'].keys())
|
||||
for key in key_list:
|
||||
del self.info['feature'][key]
|
||||
self.info = {}
|
||||
|
||||
torch.manual_seed(opts.seed)
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.manual_seed_all(opts.seed)
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
t0 = time.perf_counter()
|
||||
|
||||
try:
|
||||
from accelerate import init_empty_weights
|
||||
is_accelerate_available = True
|
||||
except:
|
||||
is_accelerate_available = False
|
||||
if not is_accelerate_available:
|
||||
self.model = model.to(self.device)
|
||||
else:
|
||||
self.model = model
|
||||
|
||||
self.z0,self.zt,self.info = self.model.inverse(inp,mask,opts)
|
||||
|
||||
# if self.offload:
|
||||
# self.model.cpu()
|
||||
# torch.cuda.empty_cache()
|
||||
|
||||
t1 = time.perf_counter()
|
||||
print(f"inversion Done in {t1 - t0:.1f}s.")
|
||||
return None
|
||||
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def edit(self,
|
||||
opts,inp_target,mask,
|
||||
):
|
||||
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
torch.manual_seed(opts.seed)
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.manual_seed_all(opts.seed)
|
||||
|
||||
torch.cuda.empty_cache()
|
||||
#self.model = self.model.to(self.device)
|
||||
x = self.model.denoise(self.z0,self.zt,inp_target,mask,opts,self.info)
|
||||
|
||||
# if self.offload:
|
||||
# self.model.cpu()
|
||||
# torch.cuda.empty_cache()
|
||||
|
||||
return x
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def encode(self,init_image, torch_device):
|
||||
init_image = torch.from_numpy(init_image).permute(2, 0, 1).float() / 127.5 - 1
|
||||
init_image = init_image.unsqueeze(0)
|
||||
init_image = init_image.to(torch_device)
|
||||
self.ae.encoder.to(torch_device)
|
||||
|
||||
init_image = self.ae.encode(init_image).to(torch.bfloat16)
|
||||
return init_image
|
||||
@@ -1,59 +0,0 @@
|
||||
|
||||
from dataclasses import dataclass
|
||||
import torch
|
||||
|
||||
|
||||
@dataclass
|
||||
class SamplingOptions:
|
||||
source_prompt: str = ''
|
||||
target_prompt: str = ''
|
||||
width: int = 1366
|
||||
height: int = 768
|
||||
inversion_num_steps: int = 0
|
||||
denoise_num_steps: int = 0
|
||||
skip_step: int = 0
|
||||
inversion_guidance: float = 1.0
|
||||
denoise_guidance: float = 1.0
|
||||
seed: int = 42
|
||||
re_init: bool = False
|
||||
attn_mask: bool = False
|
||||
attn_scale: float = 1.0
|
||||
|
||||
|
||||
class FluxEditor_kv_Wrapper_inf:
|
||||
def __init__(self, offload,device):
|
||||
|
||||
self.device = device
|
||||
self.offload = offload
|
||||
self.is_schnell = False
|
||||
#self.model = Flux_kv_edit_inf(device="cpu" if self.offload else self.device,model_path=model_path,name='flux-dev')
|
||||
#self.model.eval()
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def edit(self, model,
|
||||
opts,inp,inp_target,mask
|
||||
):
|
||||
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
torch.manual_seed(opts.seed)
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.manual_seed_all(opts.seed)
|
||||
try:
|
||||
from accelerate import init_empty_weights
|
||||
is_accelerate_available = True
|
||||
except:
|
||||
is_accelerate_available = False
|
||||
if not is_accelerate_available:
|
||||
self.model = model.to(self.device)
|
||||
else:
|
||||
self.model = model
|
||||
x = self.model(inp, inp_target, mask, opts)
|
||||
|
||||
# if self.offload:
|
||||
# self.model.cpu()
|
||||
# torch.cuda.empty_cache()
|
||||
|
||||
print("End Edit")
|
||||
return x
|
||||
-337
@@ -1,337 +0,0 @@
|
||||
# !/usr/bin/env python
|
||||
# -*- coding: UTF-8 -*-
|
||||
import os
|
||||
from PIL import Image
|
||||
import torch
|
||||
import numpy as np
|
||||
from .flux.sampling import prepare
|
||||
from einops import rearrange
|
||||
from .flux.util import load_ae_cf
|
||||
from .node_utils import cleanup,pil2narry
|
||||
from .gradio_kv_edit import FluxEditor_kv_Wrapper
|
||||
from .flux.kv_edit import Flux_kv_edit_inf,Flux_kv_edit
|
||||
from .gradio_kv_edit_inf import FluxEditor_kv_Wrapper_inf,SamplingOptions
|
||||
from comfy import model_management
|
||||
import comfy.model_management
|
||||
import folder_paths
|
||||
MAX_SEED = np.iinfo(np.int32).max
|
||||
current_node_path = os.path.dirname(os.path.abspath(__file__))
|
||||
|
||||
device = torch.device(
|
||||
"cuda:0") if torch.cuda.is_available() else torch.device(
|
||||
"mps") if torch.backends.mps.is_available() else torch.device(
|
||||
"cpu")
|
||||
|
||||
|
||||
class KV_Edit_Load:
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"unet": (["none"]+folder_paths.get_filename_list("diffusion_models"),),
|
||||
"offload":("BOOLEAN",{"default":True}),
|
||||
"use_inf":("BOOLEAN",{"default":True}),
|
||||
},
|
||||
"optional": { "model":("MODEL",),},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MODEL_KVEDIT", )
|
||||
RETURN_NAMES = ("model",)
|
||||
FUNCTION = "main"
|
||||
CATEGORY = "KV_Edit"
|
||||
|
||||
def main(self, unet,offload,use_inf,**kwargs):
|
||||
cf_model=kwargs.get("model")
|
||||
if cf_model is None:
|
||||
if unet=="none":
|
||||
raise Exception("Please select a model")
|
||||
else:
|
||||
cf_model=folder_paths.get_full_path("diffusion_models", unet)
|
||||
device_val = "cpu" if offload else device
|
||||
if not use_inf:
|
||||
model=Flux_kv_edit(device_val,cf_model,'flux-dev')
|
||||
if not isinstance(cf_model, str):
|
||||
del cf_model
|
||||
cleanup()
|
||||
pipeline=FluxEditor_kv_Wrapper(offload,device)
|
||||
else:
|
||||
model=Flux_kv_edit_inf(device_val,cf_model,'flux-dev')
|
||||
if not isinstance(cf_model, str):
|
||||
del cf_model
|
||||
cleanup()
|
||||
pipeline=FluxEditor_kv_Wrapper_inf(offload,device)
|
||||
|
||||
|
||||
|
||||
comfy.model_management.unload_all_models()
|
||||
comfy.model_management.soft_empty_cache()
|
||||
cleanup()
|
||||
return ({"pipeline":pipeline,"model":model,"use_inf":use_inf},)
|
||||
|
||||
|
||||
class KV_Edit_PreData:
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"clip":("CLIP",),
|
||||
"vae":("VAE",),
|
||||
"image": ("IMAGE", ),
|
||||
"mask":("MASK",),# B H W
|
||||
"source_prompt": ("STRING", {"default": "in a cluttered wooden cabin,a workbench holds a green neon sign that reads 'I love nana'.", "multiline": True,"tooltip": "The source_prompt to be encoded."}),
|
||||
"target_prompt": ("STRING", {"default": "in a cluttered wooden cabin,a workbench holds a green neon sign that reads 'I love here'.", "multiline": True,"tooltip": "The target_prompt to be encoded."}),
|
||||
},
|
||||
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("CONDITION_KV", )
|
||||
RETURN_NAMES = ("condition",)
|
||||
FUNCTION = "main"
|
||||
CATEGORY = "KV_Edit"
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def encode(self,init_image,ae, torch_device):
|
||||
init_image = torch.from_numpy(init_image).permute(2, 0, 1).float() / 127.5 - 1
|
||||
init_image = init_image.unsqueeze(0)
|
||||
init_image = init_image.to(torch_device).to(torch.bfloat16)
|
||||
ae.encoder.to(torch_device)
|
||||
|
||||
init_image = ae.encode(init_image).to(torch.bfloat16)
|
||||
return init_image
|
||||
|
||||
def main(self,clip, vae,image,mask,source_prompt,target_prompt,):
|
||||
ae=load_ae_cf(vae).to(device=device,dtype=torch.bfloat16)
|
||||
|
||||
# encode image
|
||||
np_image=image.squeeze().mul(255).clamp(0, 255).byte().numpy()
|
||||
shape = np_image.shape
|
||||
height = shape[0] if shape[0] % 16 == 0 else shape[0] - shape[0] % 16
|
||||
width = shape[1] if shape[1] % 16 == 0 else shape[1] - shape[1] % 16
|
||||
np_image = np_image[:height, :width, :]
|
||||
latent_image = self.encode(np_image,ae, device).to(device)
|
||||
#print(latent_image.shape) #torch.Size([1, 16, 64, 64])
|
||||
|
||||
width,height = latent_image.shape[3]*8,latent_image.shape[2]*8
|
||||
|
||||
if mask.ndim == 2:
|
||||
mask = mask.unsqueeze(0)
|
||||
|
||||
# mask crop
|
||||
mask = mask.reshape((-1, mask.shape[-2], mask.shape[-1]))
|
||||
out = mask[:, 0:height, 0:width]
|
||||
mask = out.unsqueeze(0).to(device)
|
||||
mask[mask > 0] = 1
|
||||
mask=mask.to(torch.bfloat16)
|
||||
#print(mask.shape) #torch.Size([1, 1, 512, 512])
|
||||
|
||||
with torch.no_grad():
|
||||
if clip is None:
|
||||
raise Exception("clip is None")
|
||||
inp = prepare(clip,device,latent_image, prompt=source_prompt)
|
||||
inp_target = prepare(clip,device, latent_image, prompt=target_prompt)
|
||||
|
||||
comfy.model_management.unload_all_models()
|
||||
comfy.model_management.soft_empty_cache()
|
||||
cleanup()
|
||||
|
||||
return ({"inp":inp,"inp_target":inp_target,"mask":mask,"source_prompt":source_prompt,"target_prompt":target_prompt,"size":(width,height),"ae":ae},)
|
||||
|
||||
class KV_Edit_Sampler:
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"model": ("MODEL_KVEDIT",),
|
||||
"condition": ("CONDITION_KV",),
|
||||
"seed": ("INT", {"default": 0, "min": 0, "max": MAX_SEED, "step": 1, "display": "number"}),
|
||||
"inversion_steps": ("INT", {"default": 28, "min": 1, "max": 1024, "step": 1, "display": "number"}),
|
||||
"inversion_guidance": ("FLOAT", {"default": 1.5, "min": 0.0, "max": 10.0, "step": 0.1}),
|
||||
"denoise_steps": ("INT", {"default": 28, "min": 1, "max": 1024, "step": 1, "display": "number"}),
|
||||
"denoise_guidance": ("FLOAT", {"default": 5.0, "min": 0.0, "max": 10.0, "step": 0.1}),
|
||||
"skip_step":("INT", {"default": 1, "min": 0, "max":30 ,"step": 1, "display": "number"}),
|
||||
"attn_scale":("FLOAT", {"default": 1.0, "min": 0.0, "max": 5.0, "step": 0.1}),
|
||||
"re_init":("BOOLEAN",{"default":False}),
|
||||
"attn_mask":("BOOLEAN",{"default":False}),},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE", )
|
||||
RETURN_NAMES = ("image",)
|
||||
FUNCTION = "main"
|
||||
CATEGORY = "KV_Edit"
|
||||
|
||||
def main(self,model, condition,seed,inversion_steps,inversion_guidance,denoise_steps,denoise_guidance,skip_step,attn_scale,re_init,attn_mask):
|
||||
|
||||
use_inf=model.get("use_inf")
|
||||
pipeline=model.get("pipeline")
|
||||
model_int=model.get("model")
|
||||
|
||||
inp=condition.get("inp")
|
||||
inp_target=condition.get("inp_target")
|
||||
mask=condition.get("mask")
|
||||
ae=condition.get("ae")
|
||||
|
||||
opts = SamplingOptions(
|
||||
source_prompt=condition.get("source_prompt"),
|
||||
target_prompt=condition.get("target_prompt"),
|
||||
width=condition.get("size")[0],
|
||||
height=condition.get("size")[1],
|
||||
inversion_num_steps=inversion_steps,
|
||||
denoise_num_steps=denoise_steps,
|
||||
skip_step=skip_step,
|
||||
inversion_guidance=inversion_guidance,
|
||||
denoise_guidance=denoise_guidance,
|
||||
seed=seed,
|
||||
re_init=re_init,
|
||||
attn_mask=attn_mask,
|
||||
attn_scale=attn_scale
|
||||
)
|
||||
|
||||
# inverse
|
||||
if use_inf:
|
||||
print("start edit")
|
||||
try:
|
||||
output_latent=pipeline.edit(model_int,opts,inp,inp_target,mask)
|
||||
except model_management.OOM_EXCEPTION:
|
||||
print("get OOM,try again ")
|
||||
output_latent=pipeline.edit(model_int,opts,inp,inp_target,mask)
|
||||
else:
|
||||
try:
|
||||
print("start inverse")
|
||||
if attn_mask:
|
||||
mask_inverse = mask
|
||||
else:
|
||||
mask_inverse=None
|
||||
opts.skip_step=0
|
||||
pipeline.inverse(model_int,opts,mask_inverse,inp)
|
||||
print("inverse done,start edit")
|
||||
# edit
|
||||
opts.skip_step=skip_step
|
||||
output_latent=pipeline.edit(opts,inp_target,mask)
|
||||
except model_management.OOM_EXCEPTION:
|
||||
if attn_mask:
|
||||
mask_inverse = mask
|
||||
else:
|
||||
mask_inverse=None
|
||||
opts.skip_step=0
|
||||
pipeline.inverse(model_int,opts,mask_inverse,inp)
|
||||
print("inverse done,start edit")
|
||||
# edit
|
||||
opts.skip_step=skip_step
|
||||
output_latent=pipeline.edit(opts,inp_target,mask)
|
||||
with torch.autocast(device_type=device.type, dtype=torch.bfloat16):
|
||||
x = ae.decode(output_latent.to(dtype=torch.bfloat16))
|
||||
|
||||
x = x.clamp(-1, 1)
|
||||
x = x.float().cpu()
|
||||
x = rearrange(x[0], "c h w -> h w c")
|
||||
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.synchronize()
|
||||
|
||||
img = Image.fromarray((127.5 * (x + 1.0)).cpu().byte().numpy())
|
||||
|
||||
# out = {}
|
||||
# out["samples"] = output_latent
|
||||
cleanup()
|
||||
return (pil2narry(img),)
|
||||
|
||||
class KV_Edit_Load_LoRA:
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"model": ("MODEL_KVEDIT",),
|
||||
"lora_name": (["none"] + folder_paths.get_filename_list("loras"),),
|
||||
"lora_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.01}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MODEL_KVEDIT",)
|
||||
RETURN_NAMES = ("model",)
|
||||
FUNCTION = "apply_lora"
|
||||
CATEGORY = "KV_Edit"
|
||||
|
||||
def apply_lora(self, model, lora_name, lora_scale):
|
||||
if lora_name == "none":
|
||||
return (model,)
|
||||
|
||||
try:
|
||||
# 获取LoRA文件路径
|
||||
lora_path = folder_paths.get_full_path("loras", lora_name)
|
||||
if not lora_path:
|
||||
print(f"LoRA file not found: {lora_name}")
|
||||
return (model,)
|
||||
model_int = model.get("model")
|
||||
self._apply_lora_weights(model_int, lora_path, lora_scale)
|
||||
|
||||
print(f"Successfully applied LoRA: {lora_name} with scale {lora_scale}")
|
||||
return (model,)
|
||||
|
||||
except Exception as e:
|
||||
print(f"Failed to apply LoRA {lora_name}: {str(e)}")
|
||||
return (model,)
|
||||
|
||||
def _apply_lora_weights(self, model, lora_path, scale):
|
||||
|
||||
from safetensors.torch import load_file as load_sft
|
||||
|
||||
lora_sd = load_sft(lora_path, device="cpu")
|
||||
|
||||
# 获取模型状态字典
|
||||
model_sd = model.state_dict()
|
||||
|
||||
# 应用LoRA权重
|
||||
for key in lora_sd:
|
||||
if "lora_up" in key:
|
||||
# 找到对应的down权重和原始权重
|
||||
down_key = key.replace("lora_up", "lora_down")
|
||||
original_key = self._get_original_key(key)
|
||||
|
||||
if down_key in lora_sd and original_key in model_sd:
|
||||
up_weight = lora_sd[key]
|
||||
down_weight = lora_sd[down_key]
|
||||
original_weight = model_sd[original_key]
|
||||
|
||||
# 计算LoRA增量并应用
|
||||
with torch.no_grad():
|
||||
lora_delta = (down_weight @ up_weight) * scale
|
||||
model_sd[original_key].copy_(original_weight + lora_delta)
|
||||
|
||||
# 更新模型权重
|
||||
model.load_state_dict(model_sd)
|
||||
del lora_sd
|
||||
|
||||
def _get_original_key(self, lora_key):
|
||||
"""从LoRA键名获取原始模型键名"""
|
||||
# 移除LoRA特定的后缀
|
||||
original_key = lora_key.replace(".lora_up.weight", ".weight")
|
||||
original_key = original_key.replace(".lora_down.weight", ".weight")
|
||||
return original_key
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"KV_Edit_Load": KV_Edit_Load,
|
||||
"KV_Edit_PreData": KV_Edit_PreData,
|
||||
"KV_Edit_Sampler": KV_Edit_Sampler,
|
||||
"KV_Edit_Load_LoRA": KV_Edit_Load_LoRA, # 添加LoRA节点
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"KV_Edit_Load": "KV_Edit_Load",
|
||||
"KV_Edit_PreData": "KV_Edit_PreData",
|
||||
"KV_Edit_Sampler": "KV_Edit_Sampler",
|
||||
"KV_Edit_Load_LoRA": "KV_Edit_Load_LoRA", # 添加LoRA节点
|
||||
}
|
||||
+381
@@ -0,0 +1,381 @@
|
||||
# -*- coding: UTF-8 -*-
|
||||
"""KV-Edit (https://github.com/Xilluill/KV-Edit) re-implemented on comfy's *native*
|
||||
FLUX DiT stack.
|
||||
|
||||
Upstream ships its own copy of the BFL flux model (``flux/model.py`` +
|
||||
``flux/modules/layers.py``) and mods two things:
|
||||
|
||||
1. ``DoubleStreamBlock_kv`` / ``SingleStreamBlock_kv``: during the *source* pass the
|
||||
image part of K and V is cached per block; during the *target* pass the cached
|
||||
K/V rows of the **background** are restored (while the edited region keeps the
|
||||
fresh K/V), which is what makes the background pixel-exact.
|
||||
2. two ODE drivers: an inversion-free one (``denoise_kv_inf``, upstream default) and
|
||||
a real RF-inversion one (``denoise_kv`` + ``Flux_kv_edit``).
|
||||
|
||||
This module re-implements both against comfy's own ``comfy.ldm.flux`` so the plugin
|
||||
needs **no diffusers, no transformers, no accelerate, no huggingface_hub**.
|
||||
|
||||
Why a full token sequence works here (the key porting decision)
|
||||
--------------------------------------------------------------
|
||||
Upstream feeds the target pass **only the edited tokens** (``L_mask`` rows) and, inside
|
||||
every block, splices their K/V into the cached source K/V. comfy's ``Flux._forward``
|
||||
always builds a rectangular token grid from a ``[B, C, H, W]`` latent, so we cannot feed
|
||||
a token subset. We therefore feed the **whole sequence** and, in the target phase,
|
||||
overwrite the *background* rows of K/V with the cached ones. That is exactly
|
||||
equivalent:
|
||||
|
||||
* attention is only ever read at the text rows and at the edited rows - the
|
||||
K/V tensors they see are ``[txt | cached_bg | fresh_mask]`` in both stacks;
|
||||
* the per-row projections / MLP / modulation are applied independently to each row,
|
||||
so the discarded background rows never leak into the edited rows;
|
||||
* RoPE: because upstream's target sequence is a subset, ``Flux_kv.forward`` has to
|
||||
slice the full-grid rope down to the mask positions (``info['pe_mask']``, the
|
||||
[txt | edited tokens] slice) and feed that instead of the plain ``pe``. With the
|
||||
full grid the stock RoPE is already the right position for every row, so that
|
||||
whole slice disappears;
|
||||
* the token grid (``h/16 x w/16``, row major) and the RoPE id layout are identical
|
||||
in the two stacks - no index remapping.
|
||||
|
||||
The upstream quirks that are *reproduced on purpose* (do not "fix" them):
|
||||
* the source-trajectory noise is drawn from ``Generator(device).manual_seed(0)``
|
||||
**inside** the loop -> every entry of ``noise_list`` is the same tensor, and the
|
||||
node's ``seed`` has no effect in the inversion-free path;
|
||||
* in the two-pass path the K/V cache is stored under the key of ``t_prev`` during
|
||||
inversion but looked up under ``t_curr`` during denoising -> the source K/V of one
|
||||
step is the one produced *one step later* on the schedule;
|
||||
* the saved source latent of the previous step is blended back with weight
|
||||
``(1 - mask)`` before every denoise step (a no-op for a hard mask).
|
||||
"""
|
||||
|
||||
import math
|
||||
|
||||
import torch
|
||||
|
||||
import comfy.model_management
|
||||
|
||||
PATCH = 2
|
||||
BASE_SEQ = 256
|
||||
MAX_SEQ = 4096
|
||||
BASE_SHIFT = 0.5
|
||||
MAX_SHIFT = 1.15
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------------------
|
||||
# schedule - verbatim upstream get_schedule()
|
||||
# --------------------------------------------------------------------------------------
|
||||
|
||||
def get_lin_function(x1=256, y1=0.5, x2=4096, y2=1.15):
|
||||
m = (y2 - y1) / (x2 - x1)
|
||||
b = y1 - m * x1
|
||||
return lambda x: m * x + b
|
||||
|
||||
|
||||
def time_shift(mu, sigma, t):
|
||||
return math.exp(mu) / (math.exp(mu) + (1.0 / t - 1.0) ** sigma)
|
||||
|
||||
|
||||
def get_schedule(num_steps, image_seq_len, shift=True):
|
||||
"""Upstream ``flux.sampling.get_schedule``.
|
||||
|
||||
Note that ``linspace(1, 0, n+1)`` + shift is bit-identical to diffusers'
|
||||
``linspace(1, 1/n, n)`` + shift + terminal 0, i.e. this *is* the official FLUX
|
||||
sigma table - with a **resolution dependent** mu, which comfy's
|
||||
``ModelSamplingFlux`` (hardcoded shift=1.15) is not. So we drive it ourselves.
|
||||
"""
|
||||
timesteps = torch.linspace(1, 0, num_steps + 1)
|
||||
if shift:
|
||||
mu = get_lin_function(y1=BASE_SHIFT, y2=MAX_SHIFT)(image_seq_len)
|
||||
timesteps = time_shift(mu, 1.0, timesteps)
|
||||
return timesteps
|
||||
|
||||
|
||||
def describe(steps, tokens):
|
||||
mu = get_lin_function(y1=BASE_SHIFT, y2=MAX_SHIFT)(tokens)
|
||||
return f"steps={steps} tokens={tokens} mu={mu:.4f} shift={math.exp(mu):.3f}"
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------------------
|
||||
# token <-> latent
|
||||
# --------------------------------------------------------------------------------------
|
||||
|
||||
def pack_tokens(x, patch_size=PATCH):
|
||||
"""[B, C, H, W] -> [B, L, C*p*p], the inverse of comfy's rearrange in Flux._forward."""
|
||||
b, c, h, w = x.shape
|
||||
p = patch_size
|
||||
return (x.view(b, c, h // p, p, w // p, p)
|
||||
.permute(0, 2, 4, 1, 3, 5)
|
||||
.reshape(b, (h // p) * (w // p), c * p * p))
|
||||
|
||||
|
||||
def unpack_tokens(t, h, w, patch_size=PATCH):
|
||||
"""[B, L, C*p*p] -> [B, C, H, W]."""
|
||||
b, l, d = t.shape
|
||||
p = patch_size
|
||||
c = d // (p * p)
|
||||
return (t.view(b, h // p, w // p, c, p, p)
|
||||
.permute(0, 3, 1, 4, 2, 5)
|
||||
.reshape(b, c, h, w))
|
||||
|
||||
|
||||
def token_grid(latent_shape, patch_size=PATCH):
|
||||
h = int(latent_shape[-2])
|
||||
w = int(latent_shape[-1])
|
||||
return ((h + patch_size // 2) // patch_size, (w + patch_size // 2) // patch_size)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------------------
|
||||
# mask -> token indices (upstream semantics: bilinear -> >0 -> 16ch -> patchify)
|
||||
# --------------------------------------------------------------------------------------
|
||||
|
||||
def mask_to_tokens(mask, h_lat, w_lat, patch_size=PATCH):
|
||||
"""comfy MASK [B,H,W] in [0,1] -> (edited_idx, background_idx, token_mask[1,L,C*p*p]).
|
||||
|
||||
Upstream resizes the mask with **bilinear** to the latent resolution and keeps every
|
||||
non-zero pixel (``mask[mask > 0] = 1``), then repeats it to 16 channels and
|
||||
patchifies. A token counts as edited when at least one of its 4 sub-pixels is > 0.
|
||||
The per-channel token mask is kept because upstream scales the velocity update with
|
||||
it channel-wise.
|
||||
"""
|
||||
m = mask
|
||||
if m is None:
|
||||
raise ValueError("mask is required")
|
||||
if m.ndim == 4: # [B, 1, H, W] is accepted too
|
||||
m = m[:, 0]
|
||||
if m.ndim == 2:
|
||||
m = m.unsqueeze(0)
|
||||
m = m[0].detach().float().cpu()
|
||||
# crop to what the latent covers (upstream crops the image to a multiple of 16)
|
||||
m = m[: h_lat * 8, : w_lat * 8]
|
||||
m = (m > 0).float().unsqueeze(0).unsqueeze(0)
|
||||
m = torch.nn.functional.interpolate(m, size=(h_lat, w_lat), mode="bilinear",
|
||||
align_corners=False)
|
||||
m = (m > 0).float()
|
||||
m = m.repeat(1, 16, 1, 1)
|
||||
tok = pack_tokens(m, patch_size) # [1, L, 64]
|
||||
keep = tok.sum(-1)[0] > 0.5 # [L]
|
||||
edited = torch.nonzero(keep, as_tuple=False).reshape(-1)
|
||||
bg = torch.nonzero(~keep, as_tuple=False).reshape(-1)
|
||||
return edited, bg, tok
|
||||
|
||||
|
||||
def pixel_mask_from_tokens(tok, h_tok, w_tok, patch_size=PATCH):
|
||||
"""[h,w] float mask, 1 on the pixels a *edited* token owns.
|
||||
|
||||
Used to composite the result back onto the input latent, so the untouched area
|
||||
stays **bit-identical** instead of surviving a raw->scaled->raw round trip
|
||||
(that round trip costs 1 fp32 ulp, ~9.5e-07 at the magnitude flux latents live
|
||||
at, and it would accumulate over chained edits).
|
||||
"""
|
||||
keep = (tok[0].sum(-1) > 0.5).view(h_tok, w_tok).float()
|
||||
return torch.kron(keep, torch.ones((patch_size, patch_size), device=keep.device,
|
||||
dtype=keep.dtype))
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------------------
|
||||
# the attention patch
|
||||
# --------------------------------------------------------------------------------------
|
||||
|
||||
class KVState:
|
||||
"""Mutable per-run state read by the ``attn1_patch``.
|
||||
|
||||
``phase`` 0 = source pass (cache the image K/V), 1 = target pass (restore the
|
||||
background rows of K/V from the cache).
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.enabled = False
|
||||
self.phase = 0
|
||||
self.store = {} # (block_type, block_index) -> (k_img, v_img)
|
||||
self.vault = {} # timestep key -> store (two-pass mode only)
|
||||
self.edited_idx = None # token indices of the edited region
|
||||
self.bg_idx = None # token indices that keep the source K/V
|
||||
self.use_attn_mask = False
|
||||
self.attn_scale = None # None or a float
|
||||
self.kv_device = None # where the cached K/V live
|
||||
self.n_fire = 0
|
||||
self._cache = {}
|
||||
|
||||
def reset(self):
|
||||
self.store = {}
|
||||
self.vault = {}
|
||||
self.n_fire = 0
|
||||
self._cache = {}
|
||||
|
||||
def attn_mask_for(self, txt_len, n_img, device):
|
||||
key = ("mask", txt_len, n_img)
|
||||
hit = self._cache.get(key)
|
||||
if hit is None:
|
||||
s = txt_len + n_img
|
||||
am = torch.zeros(s, s, dtype=torch.bool, device=device)
|
||||
am[:txt_len, :] = True
|
||||
t = torch.arange(txt_len, device=device)
|
||||
if self.edited_idx is not None and self.edited_idx.numel():
|
||||
si = txt_len + self.edited_idx
|
||||
am[si.unsqueeze(1), si.unsqueeze(0)] = True
|
||||
am[si.unsqueeze(1), t.unsqueeze(0)] = True
|
||||
if self.bg_idx is not None and self.bg_idx.numel():
|
||||
sb = txt_len + self.bg_idx
|
||||
am[sb.unsqueeze(1), sb.unsqueeze(0)] = True
|
||||
am[sb.unsqueeze(1), t.unsqueeze(0)] = True
|
||||
hit = am.unsqueeze(0)
|
||||
self._cache[key] = hit
|
||||
return hit
|
||||
|
||||
def attn_scale_for(self, txt_len, n_img, dtype, device):
|
||||
key = ("scale", txt_len, n_img, dtype)
|
||||
hit = self._cache.get(key)
|
||||
if hit is None:
|
||||
s = txt_len + n_img
|
||||
a = torch.zeros(1, 1, s, dtype=dtype, device=device)
|
||||
if self.bg_idx is not None and self.bg_idx.numel():
|
||||
a[0, 0, txt_len + self.bg_idx] = float(self.attn_scale)
|
||||
hit = a
|
||||
self._cache[key] = hit
|
||||
return hit
|
||||
|
||||
|
||||
def make_attn_patch(state: KVState):
|
||||
"""comfy ``transformer_options["patches"]["attn1_patch"]`` callable.
|
||||
|
||||
Called by every flux block (19 double + 38 single) with q/k/v of shape
|
||||
``[B, heads, seq, head_dim]``; the sequence is ``[text | image]`` in both stacks.
|
||||
"""
|
||||
|
||||
def patch(q, k, v, pe=None, attn_mask=None, extra_options=None):
|
||||
if not state.enabled:
|
||||
return {}
|
||||
eo = extra_options or {}
|
||||
img_slice = eo.get("img_slice", None)
|
||||
if img_slice is None:
|
||||
return {}
|
||||
txt_len = int(img_slice[0])
|
||||
n_img = int(k.shape[2]) - txt_len
|
||||
blk = (eo.get("block_type", "?"), int(eo.get("block_index", 0)))
|
||||
out = {}
|
||||
|
||||
if state.phase == 0:
|
||||
# ---- source pass: remember the image rows of K/V ----------------------
|
||||
k_img = k[:, :, txt_len:, :]
|
||||
v_img = v[:, :, txt_len:, :]
|
||||
if state.kv_device is not None and k.device != state.kv_device:
|
||||
k_img = k_img.to(state.kv_device)
|
||||
v_img = v_img.to(state.kv_device)
|
||||
else:
|
||||
# clone so the (much bigger) text part of the cat() can be freed
|
||||
k_img = k_img.clone()
|
||||
v_img = v_img.clone()
|
||||
state.store[blk] = (k_img, v_img)
|
||||
if state.use_attn_mask and state.bg_idx is not None and state.bg_idx.numel():
|
||||
out["attn_mask"] = state.attn_mask_for(txt_len, n_img, q.device)
|
||||
else:
|
||||
# ---- target pass: put the cached background rows back -----------------
|
||||
entry = state.store.get(blk)
|
||||
if entry is not None and state.bg_idx is not None and state.bg_idx.numel():
|
||||
src_k, src_v = entry
|
||||
dev = k.device
|
||||
# select the background rows *where the cache lives*: with a CPU cache
|
||||
# only the rows we actually need cross the bus (a 1024x1024 image is
|
||||
# 2.9 GB of K/V per step - too much for a 12 GB card).
|
||||
bi = state.bg_idx.to(src_k.device)
|
||||
sel_k = src_k.index_select(2, bi).to(dev)
|
||||
sel_v = src_v.index_select(2, bi).to(dev)
|
||||
idx = txt_len + state.bg_idx.to(dev)
|
||||
k = k.index_copy(2, idx, sel_k)
|
||||
v = v.index_copy(2, idx, sel_v)
|
||||
state.n_fire += 1
|
||||
out["k"] = k
|
||||
out["v"] = v
|
||||
if state.attn_scale:
|
||||
out["attn_mask"] = state.attn_scale_for(txt_len, n_img, q.dtype, q.device)
|
||||
return out
|
||||
|
||||
return patch
|
||||
|
||||
|
||||
def build_transformer_options(patcher, state: KVState):
|
||||
"""The patcher's own options + our patch (so other hooks keep working)."""
|
||||
base = {}
|
||||
model_options = getattr(patcher, "model_options", None) or {}
|
||||
own = model_options.get("transformer_options", None) or {}
|
||||
if own:
|
||||
base = dict(own)
|
||||
patches = dict(base.get("patches", {}) or {})
|
||||
patches["attn1_patch"] = list(patches.get("attn1_patch", []) or []) + [make_attn_patch(state)]
|
||||
base["patches"] = patches
|
||||
return base
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------------------
|
||||
# raw velocity proxy (comfy's calculate_denoised returns x0 = x - sigma*v)
|
||||
# --------------------------------------------------------------------------------------
|
||||
|
||||
class RawVelocitySampling(torch.nn.Module):
|
||||
def __init__(self, base):
|
||||
super().__init__()
|
||||
self.base = base
|
||||
|
||||
def calculate_input(self, sigma, noise):
|
||||
return self.base.calculate_input(sigma, noise)
|
||||
|
||||
def timestep(self, sigma):
|
||||
return self.base.timestep(sigma)
|
||||
|
||||
def calculate_denoised(self, sigma, model_output, model_input):
|
||||
return model_output.float()
|
||||
|
||||
def percent_to_sigma(self, percent):
|
||||
return self.base.percent_to_sigma(percent)
|
||||
|
||||
def sigma(self, timestep):
|
||||
return self.base.sigma(timestep)
|
||||
|
||||
@property
|
||||
def sigma_min(self):
|
||||
return self.base.sigma_min
|
||||
|
||||
@property
|
||||
def sigma_max(self):
|
||||
return self.base.sigma_max
|
||||
|
||||
|
||||
class raw_velocity:
|
||||
"""Context manager: swap the model's model_sampling for the raw-velocity proxy."""
|
||||
|
||||
def __init__(self, base_model):
|
||||
self.base_model = base_model
|
||||
self.orig = None
|
||||
|
||||
def __enter__(self):
|
||||
self.orig = self.base_model.model_sampling
|
||||
self.base_model.model_sampling = RawVelocitySampling(self.orig)
|
||||
return self.base_model
|
||||
|
||||
def __exit__(self, *args):
|
||||
self.base_model.model_sampling = self.orig
|
||||
self.orig = None
|
||||
return False
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------------------
|
||||
# conditioning helpers
|
||||
# --------------------------------------------------------------------------------------
|
||||
|
||||
def _cond_rows(conditioning):
|
||||
"""comfy CONDITIONING ([[emb, dict], ...]) -> a cond row dict for extra_conds."""
|
||||
if conditioning is None or len(conditioning) == 0:
|
||||
raise ValueError("empty conditioning input")
|
||||
emb, extra = conditioning[0][0], conditioning[0][1]
|
||||
row = dict(extra)
|
||||
row["cross_attn"] = emb
|
||||
return row
|
||||
|
||||
|
||||
def rows_to_model_conds(base_model, row, noise, device):
|
||||
d = dict(row)
|
||||
d.setdefault("noise", noise)
|
||||
d.setdefault("device", device)
|
||||
return comfy.samplers.cond_cat([base_model.extra_conds(**d)])
|
||||
|
||||
|
||||
def kv_bytes(n_img, n_blocks, heads=24, head_dim=128, itemsize=2):
|
||||
return n_blocks * 2 * heads * n_img * head_dim * itemsize
|
||||
+262
@@ -0,0 +1,262 @@
|
||||
# -*- coding: UTF-8 -*-
|
||||
"""ComfyUI nodes for KV-Edit (native FLUX DiT stack).
|
||||
|
||||
The plugin deliberately owns **no** CLIP, no VAE and no sampler of its own:
|
||||
|
||||
Load Diffusion Model -> MODEL ─────────────────────────────────┐
|
||||
Load CLIP -> CLIPTextEncode(source prompt) -> CONDITIONING ────┤
|
||||
Load CLIP -> CLIPTextEncode(target prompt) -> CONDITIONING ────┤
|
||||
Load Image -> VAEEncode -> LATENT (the source image) ──────────┤
|
||||
Load Image -> MaskEditor / solid mask -> MASK ────────────┤
|
||||
▼
|
||||
KV_Edit_Sample -> LATENT
|
||||
▼
|
||||
VAEDecode -> IMAGE
|
||||
|
||||
Because CLIP / VAE / DiT live in different nodes, comfy can unload each of them when
|
||||
the next one needs the VRAM - the old wrapper materialised a second copy of the flux
|
||||
weights (accelerate load_checkpoint_and_dispatch on top of the comfy model) which is
|
||||
what blew up a 12 GB card and made the plugin depend on accelerate + huggingface_hub.
|
||||
|
||||
All LATENT inputs/outputs follow the stock convention: **raw (unscaled) latents**, the
|
||||
same space VAEEncode / KSampler / VAEDecode use; the scaling to the DiT's space happens
|
||||
inside the node via ``model.process_latent_in/out``.
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
import torch
|
||||
|
||||
import comfy.model_management
|
||||
import comfy.utils
|
||||
|
||||
from . import kvedit_core as core
|
||||
from . import kvedit_sampling as sampling
|
||||
|
||||
|
||||
def _activation_budget(patcher, latent_shape):
|
||||
"""VRAM to keep free for one forward (comfy's own estimate is tuned for B=1).
|
||||
|
||||
KV-Edit runs the DiT twice per step and additionally holds one step of K/V, so we
|
||||
scale comfy's formula and keep a 2 GB floor.
|
||||
"""
|
||||
base_model = patcher.model
|
||||
batch = max(1, int(latent_shape[0]))
|
||||
area = batch * int(latent_shape[-2]) * int(latent_shape[-1])
|
||||
dtype_size = comfy.model_management.dtype_size(base_model.get_dtype_inference())
|
||||
factor = getattr(base_model, "memory_usage_factor", 3.1) or 3.1
|
||||
est = int(area * dtype_size * 0.01 * factor * (1024 * 1024))
|
||||
floor = int(float(os.environ.get("KVEDIT_VRAM_BUDGET", "2")) * (1024 ** 3))
|
||||
return max(est * 2, floor)
|
||||
|
||||
|
||||
def _prepare_model(patcher, latent_shape):
|
||||
need = _activation_budget(patcher, latent_shape)
|
||||
comfy.model_management.unload_all_models()
|
||||
comfy.model_management.soft_empty_cache()
|
||||
comfy.model_management.load_models_gpu([patcher], memory_required=need,
|
||||
minimum_memory_required=need)
|
||||
try:
|
||||
patcher.pre_run()
|
||||
except Exception:
|
||||
pass
|
||||
return patcher.model, patcher.load_device
|
||||
|
||||
|
||||
def _make_apply(base_model, transformer_options, model_conds):
|
||||
def apply_fn(x, t):
|
||||
t_tensor = torch.full((x.shape[0],), float(t), device=x.device,
|
||||
dtype=torch.float32)
|
||||
return base_model.apply_model(x, t_tensor,
|
||||
transformer_options=transformer_options,
|
||||
**model_conds).float()
|
||||
return apply_fn
|
||||
|
||||
|
||||
def _block_count(base_model):
|
||||
dm = base_model.diffusion_model
|
||||
return len(dm.double_blocks) + len(dm.single_blocks)
|
||||
|
||||
|
||||
def _set_guidance(conditioning, guidance):
|
||||
"""CONDITIONING with the flux guidance value forced (FluxGuidance node optional)."""
|
||||
if conditioning is None or len(conditioning) == 0:
|
||||
raise ValueError("empty conditioning")
|
||||
out = []
|
||||
for c in conditioning:
|
||||
emb, extra = c[0], c[1]
|
||||
d = dict(extra)
|
||||
d["guidance"] = float(guidance)
|
||||
out.append([emb, d])
|
||||
return out
|
||||
|
||||
|
||||
class KV_Edit_Sample:
|
||||
"""Training-free image editing with precise background preservation (KV-Edit)."""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"model": ("MODEL", {"tooltip": "A native flux MODEL (UNETLoader / "
|
||||
"CheckpointLoaderSimple / + LoraLoader). "
|
||||
"flux1-dev-fp8 works."}),
|
||||
"source_conditioning": ("CONDITIONING", {"tooltip": "CLIPTextEncode of the "
|
||||
"prompt that describes the INPUT image."}),
|
||||
"target_conditioning": ("CONDITIONING", {"tooltip": "CLIPTextEncode of the "
|
||||
"prompt that describes the DESIRED image."}),
|
||||
"latent_image": ("LATENT", {"tooltip": "Raw latent of the source image "
|
||||
"(VAEEncode)."}),
|
||||
"mask": ("MASK", {"tooltip": "Region to EDIT (white = edited). The "
|
||||
"complement is preserved pixel-exact."}),
|
||||
"steps": ("INT", {"default": 28, "min": 1, "max": 200, "step": 1}),
|
||||
"skip_step": ("INT", {"default": 1, "min": 0, "max": 30, "step": 1,
|
||||
"tooltip": "Upstream default 1: drop the noisiest "
|
||||
"step of the schedule."}),
|
||||
"attn_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 5.0, "step": 0.1,
|
||||
"tooltip": "Additive bias on the background keys "
|
||||
"during the target pass. 0 disables it. "
|
||||
"1.0 = upstream default."}),
|
||||
"attn_mask": ("BOOLEAN", {"default": False,
|
||||
"tooltip": "Upstream `attn_mask`: hard-block "
|
||||
"mask<->background attention during the "
|
||||
"source pass. Off upstream-default; "
|
||||
"it forces a slow attention backend."}),
|
||||
"use_inf": ("BOOLEAN", {"default": True,
|
||||
"tooltip": "Inversion-free path (upstream default). "
|
||||
"OFF = a real RF inversion pass first."}),
|
||||
"source_guidance": ("FLOAT", {"default": 1.5, "min": 0.0, "max": 30.0,
|
||||
"step": 0.1}),
|
||||
"target_guidance": ("FLOAT", {"default": 5.0, "min": 0.0, "max": 30.0,
|
||||
"step": 0.1}),
|
||||
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff,
|
||||
"display": "number",
|
||||
"tooltip": "Only used by re_init on the two-pass path; "
|
||||
"upstream fixes the source-trajectory noise to "
|
||||
"seed 0 on the inversion-free path."}),
|
||||
},
|
||||
"optional": {
|
||||
"re_init": ("BOOLEAN", {"default": False,
|
||||
"tooltip": "Two-pass path only: restart the edited "
|
||||
"region from fresh noise instead of the "
|
||||
"inverted latent."}),
|
||||
"kv_offload": ("BOOLEAN", {"default": True,
|
||||
"tooltip": "Keep the per-step K/V cache on the CPU "
|
||||
"(~0.7 GB at 512, ~2.9 GB at 1024). "
|
||||
"On by default: on a 12 GB card the GPU "
|
||||
"copy does not fit next to the model."}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("LATENT",)
|
||||
RETURN_NAMES = ("latent",)
|
||||
FUNCTION = "run"
|
||||
CATEGORY = "KV_Edit"
|
||||
DESCRIPTION = (
|
||||
"KV-Edit (https://github.com/Xilluill/KV-Edit) on comfy's native FLUX DiT.\n"
|
||||
"The source pass caches the K/V of the background for every block; the target pass "
|
||||
"restores them, so the unmasked area survives the edit exactly and only the masked "
|
||||
"region follows the target prompt.\n"
|
||||
"Wire: VAEEncode(source image) -> latent_image; MASK = the region to edit -> mask; "
|
||||
"then VAEDecode the output. \n"
|
||||
"Guidance: set it here (source/target separately) OR with a FluxGuidance node - this "
|
||||
"node's values win. With an 8-step turbo LoRA use steps=8 and the flux sigma table "
|
||||
"is still driven by the resolution-dependent shift (comfy's default 1.15 is not)."
|
||||
)
|
||||
|
||||
def run(self, model, source_conditioning, target_conditioning, latent_image, mask,
|
||||
steps, skip_step, attn_scale, attn_mask, use_inf, source_guidance,
|
||||
target_guidance, seed, re_init=False, kv_offload=True):
|
||||
|
||||
raw = latent_image["samples"]
|
||||
if raw.shape[0] != 1:
|
||||
raise ValueError("KV-Edit works on one image at a time (batch must be 1)")
|
||||
device = model.load_device
|
||||
raw = raw.to(device=device, dtype=torch.float32)
|
||||
scaled = model.model.process_latent_in(raw)
|
||||
|
||||
h, w = int(scaled.shape[-2]), int(scaled.shape[-1])
|
||||
h_tok, w_tok = core.token_grid(scaled.shape)
|
||||
n_img = h_tok * w_tok
|
||||
|
||||
edited, bg, tok = core.mask_to_tokens(mask, h, w)
|
||||
if edited.numel() == 0:
|
||||
raise ValueError("the mask selects no token - nothing to edit")
|
||||
if bg.numel() == 0:
|
||||
print("[KV_Edit] warning: the mask covers the whole image, nothing is preserved")
|
||||
|
||||
sigmas = core.get_schedule(int(steps), n_img).to(device)
|
||||
|
||||
state = core.KVState()
|
||||
state.edited_idx = edited.to(device)
|
||||
state.bg_idx = bg.to(device)
|
||||
state.tok_mask = tok.to(device)
|
||||
state.use_attn_mask = bool(attn_mask)
|
||||
state.attn_scale = float(attn_scale) if float(attn_scale) != 0.0 else None
|
||||
# 57 blocks x K,V x 24 heads x L x 128 x 2 B: ~0.7 GB/step at 512, ~2.9 GB at 1024
|
||||
state.kv_device = torch.device("cpu") if kv_offload else None
|
||||
state.enabled = True
|
||||
|
||||
n_blocks = _block_count(model.model)
|
||||
kv_mem = core.kv_bytes(n_img, n_blocks)
|
||||
print(f"[KV_Edit] {h_tok}x{w_tok} tokens | {core.describe(steps, n_img)} | "
|
||||
f"edited tokens={edited.numel()} ({100.0 * edited.numel() / n_img:.1f}%) | "
|
||||
f"K/V cache {kv_mem / 2**30:.2f}GB/step x {n_blocks} blocks")
|
||||
|
||||
transformer_options = core.build_transformer_options(model, state)
|
||||
base_model, device = _prepare_model(model, scaled.shape)
|
||||
|
||||
noise_shape = torch.zeros_like(scaled)
|
||||
conds_src = core.rows_to_model_conds(
|
||||
base_model, core._cond_rows(_set_guidance(source_conditioning, source_guidance)),
|
||||
noise_shape, device)
|
||||
conds_tgt = core.rows_to_model_conds(
|
||||
base_model, core._cond_rows(_set_guidance(target_conditioning, target_guidance)),
|
||||
noise_shape, device)
|
||||
apply_src = _make_apply(base_model, transformer_options, conds_src)
|
||||
apply_tgt = _make_apply(base_model, transformer_options, conds_tgt)
|
||||
|
||||
z0_tok = core.pack_tokens(scaled, core.PATCH)
|
||||
pbar = comfy.utils.ProgressBar(max(1, int(steps)))
|
||||
|
||||
with comfy.model_management.cuda_device_context(device):
|
||||
with core.raw_velocity(base_model):
|
||||
if use_inf:
|
||||
ts = sigmas[int(skip_step):]
|
||||
z_fe = sampling.kv_edit_inversion_free(
|
||||
apply_src, apply_tgt, z0_tok, h, w, ts, state,
|
||||
progress=lambda i, n: pbar.update_absolute(i, n))
|
||||
else:
|
||||
saved = {}
|
||||
sampling.kv_invert(apply_src, z0_tok, h, w, sigmas, state, saved,
|
||||
offload=bool(kv_offload),
|
||||
progress=lambda i, n: pbar.update_absolute(i, n))
|
||||
z_fe = sampling.kv_denoise(
|
||||
apply_tgt, z0_tok, h, w, sigmas, skip_step, state, saved,
|
||||
n_inv=len(sigmas) - 1, re_init=bool(re_init), seed=seed,
|
||||
progress=lambda i, n: pbar.update_absolute(i, n))
|
||||
saved.clear()
|
||||
|
||||
out_tok = z0_tok.clone()
|
||||
out_tok[:, state.edited_idx, :] = z_fe
|
||||
out = model.model.process_latent_out(core.unpack_tokens(out_tok, h, w, core.PATCH))
|
||||
# composite onto the input latent: the untouched area is then *bit-identical*,
|
||||
# not merely "a raw->scaled->raw round trip of it" (that costs 1 fp32 ulp).
|
||||
pm = core.pixel_mask_from_tokens(tok, h_tok, w_tok).to(device=out.device,
|
||||
dtype=out.dtype)
|
||||
out = raw * (1.0 - pm) + out * pm
|
||||
|
||||
print(f"[KV_Edit] done | patch fired {state.n_fire}x")
|
||||
state.reset()
|
||||
state.enabled = False
|
||||
del state
|
||||
return ({"samples": out.to(device="cpu")},)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"KV_Edit_Sample": KV_Edit_Sample,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"KV_Edit_Sample": "KV_Edit_Sample (native DiT)",
|
||||
}
|
||||
@@ -0,0 +1,154 @@
|
||||
# -*- coding: UTF-8 -*-
|
||||
"""The two KV-Edit ODE drivers, replayed on comfy's native FLUX stack.
|
||||
|
||||
Both take the *scaled* clean latent in token space (``[1, L, 64]``), exactly like the
|
||||
upstream ``flux.sampling.denoise_kv*`` helpers, and return the edited tokens.
|
||||
"""
|
||||
|
||||
import torch
|
||||
|
||||
from .kvedit_core import PATCH, pack_tokens, unpack_tokens
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------------------
|
||||
# 1) inversion-free (upstream ``denoise_kv_inf``, the upstream default)
|
||||
# --------------------------------------------------------------------------------------
|
||||
|
||||
def kv_edit_inversion_free(apply_src, apply_tgt, z0_tok, h, w, sigmas, state,
|
||||
progress=None):
|
||||
"""Upstream ``Flux_kv_edit_inf.forward`` + ``denoise_kv_inf``.
|
||||
|
||||
One loop, two model calls per step:
|
||||
z_src = (1 - t) * z0 + t * noise (analytic source trajectory)
|
||||
z_tar = z_src[edited] - z0[edited] + z_fe (the edited region's own trajectory)
|
||||
v_fe = v_tar - v_src[edited] (velocity of the edit)
|
||||
z_fe += (t_prev - t_curr) * v_fe * mask
|
||||
"""
|
||||
device = z0_tok.device
|
||||
edited = state.edited_idx
|
||||
tok_mask = state.tok_mask[:, edited, :]
|
||||
n = len(sigmas)
|
||||
steps = max(n - 1, 0)
|
||||
|
||||
# Upstream creates a fresh ``Generator(device).manual_seed(0)`` *inside* the loop, so
|
||||
# every entry of noise_list is the same tensor. Reproduced verbatim (one draw).
|
||||
gen = torch.Generator(device=device).manual_seed(0)
|
||||
noise = torch.randn(z0_tok.size(), dtype=z0_tok.dtype, layout=z0_tok.layout,
|
||||
device=device, generator=gen)
|
||||
|
||||
z_fe = z0_tok[:, edited, :].clone()
|
||||
|
||||
for i in range(steps):
|
||||
t_curr = float(sigmas[i])
|
||||
t_prev = float(sigmas[i + 1])
|
||||
|
||||
z_src = (1.0 - t_curr) * z0_tok + t_curr * noise
|
||||
z_tar = z_src[:, edited, :] - z0_tok[:, edited, :] + z_fe
|
||||
# full token grid: the edited rows carry the target trajectory, the background
|
||||
# rows are irrelevant (their K/V is replaced by the cache inside every block).
|
||||
full_tar = z_src.clone()
|
||||
full_tar[:, edited, :] = z_tar
|
||||
|
||||
state.store = {}
|
||||
state.phase = 0
|
||||
v_src = apply_src(unpack_tokens(z_src, h, w, PATCH), t_curr)
|
||||
|
||||
state.phase = 1
|
||||
v_tar = apply_tgt(unpack_tokens(full_tar, h, w, PATCH), t_curr)
|
||||
|
||||
v_fe = (pack_tokens(v_tar, PATCH)[:, edited, :]
|
||||
- pack_tokens(v_src, PATCH)[:, edited, :])
|
||||
z_fe = z_fe + (t_prev - t_curr) * v_fe * tok_mask
|
||||
|
||||
state.store = {}
|
||||
if progress is not None:
|
||||
progress(i + 1, steps)
|
||||
|
||||
return z_fe
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------------------
|
||||
# 2) real RF inversion + denoise (upstream ``denoise_kv``)
|
||||
# --------------------------------------------------------------------------------------
|
||||
|
||||
def kv_invert(apply_src, z0_tok, h, w, sigmas, state, saved, offload=False,
|
||||
progress=None):
|
||||
"""Upstream ``Flux_kv_edit.inverse``.
|
||||
|
||||
Walks the schedule from t=0 to t=1 adding the velocity (the naive RF inversion),
|
||||
and stores, for every step, the K/V of every block *and* the pre-update latent.
|
||||
Upstream keys the store with ``t_prev``; we keep integer keys and reuse them.
|
||||
"""
|
||||
ts = sigmas.flip(0) # ascending: [0.0, s27, ..., s1, 1.0]
|
||||
z = z0_tok
|
||||
n = len(ts) - 1
|
||||
for i in range(n):
|
||||
t_curr = float(ts[i])
|
||||
t_prev = float(ts[i + 1])
|
||||
state.store = {}
|
||||
saved[i] = {"z": z.detach().clone(), "kv": state.store}
|
||||
state.phase = 0
|
||||
v = apply_src(unpack_tokens(z, h, w, PATCH), t_curr)
|
||||
if offload:
|
||||
for blk in list(state.store.keys()):
|
||||
k_img, v_img = state.store[blk]
|
||||
state.store[blk] = (k_img.to("cpu", non_blocking=False),
|
||||
v_img.to("cpu", non_blocking=False))
|
||||
saved[i]["z"] = saved[i]["z"].to("cpu", non_blocking=False)
|
||||
z = z + (t_prev - t_curr) * pack_tokens(v, PATCH)
|
||||
if progress is not None:
|
||||
progress(i + 1, n)
|
||||
return z
|
||||
|
||||
|
||||
def kv_denoise(apply_tgt, z0_tok, h, w, sigmas_full, skip, state, saved, n_inv,
|
||||
re_init=False, seed=0, progress=None):
|
||||
"""Upstream ``Flux_kv_edit.denoise``.
|
||||
|
||||
Upstream keys the inversion store with ``t_prev`` and looks it up with ``t_curr``.
|
||||
Resolving that: denoise iteration j reads the entry produced by inversion iteration
|
||||
``n_inv - 1 - skip - j``, i.e. the source K/V always come from the state **one step
|
||||
later** on the schedule. That off-by-one is reproduced on purpose.
|
||||
"""
|
||||
edited = state.edited_idx
|
||||
tok_mask = state.tok_mask[:, edited, :]
|
||||
skip = int(skip)
|
||||
ts = sigmas_full[skip:]
|
||||
n = len(ts) - 1
|
||||
|
||||
if re_init:
|
||||
# upstream: noise = randn_like(zt); zt = z0*(1-t) + noise*t ; t = ts[0]
|
||||
gen = torch.Generator(device=z0_tok.device).manual_seed(int(seed) & 0xFFFFFFFF)
|
||||
noise = torch.randn(z0_tok.size(), dtype=z0_tok.dtype, layout=z0_tok.layout,
|
||||
device=z0_tok.device, generator=gen)
|
||||
full = (z0_tok * (1.0 - float(ts[0])) + noise * float(ts[0])).contiguous()
|
||||
else:
|
||||
entry = saved.get(n_inv - 1 - skip)
|
||||
if entry is None:
|
||||
raise RuntimeError("empty inversion: nothing was stored")
|
||||
full = entry["z"].to(z0_tok.device).clone()
|
||||
|
||||
for j in range(n):
|
||||
t_curr = float(ts[j])
|
||||
t_prev = float(ts[j + 1])
|
||||
entry = saved.get(n_inv - 1 - skip - j)
|
||||
if entry is None:
|
||||
raise RuntimeError(f"missing inversion cache for step {j}")
|
||||
source = entry["z"].to(z0_tok.device)
|
||||
# upstream blends the previous source latent back with weight (1 - mask);
|
||||
# a no-op for a hard mask, kept for bit-fidelity with soft masks.
|
||||
full[:, edited, :] = (source[:, edited, :] * (1.0 - tok_mask)
|
||||
+ full[:, edited, :] * tok_mask)
|
||||
|
||||
state.store = entry["kv"]
|
||||
state.phase = 1
|
||||
v = apply_tgt(unpack_tokens(full, h, w, PATCH), t_curr)
|
||||
|
||||
full = full + (t_prev - t_curr) * pack_tokens(v, PATCH)
|
||||
if progress is not None:
|
||||
progress(j + 1, n)
|
||||
|
||||
out = z0_tok.clone()
|
||||
out[:, edited, :] = (z0_tok[:, edited, :] * (1.0 - tok_mask)
|
||||
+ full[:, edited, :] * tok_mask)
|
||||
return out[:, edited, :]
|
||||
-239
@@ -1,239 +0,0 @@
|
||||
# !/usr/bin/env python
|
||||
# -*- coding: UTF-8 -*-
|
||||
import os
|
||||
import torch
|
||||
from PIL import Image
|
||||
import numpy as np
|
||||
import cv2
|
||||
import gc
|
||||
import torch
|
||||
from comfy.utils import common_upscale,ProgressBar
|
||||
from .flux.sampling import prepare
|
||||
from huggingface_hub import hf_hub_download
|
||||
|
||||
cur_path = os.path.dirname(os.path.abspath(__file__))
|
||||
device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
|
||||
|
||||
def latent_crop( samples, width, height, x, y):
|
||||
s = samples.copy()
|
||||
samples = samples['samples']
|
||||
x = x // 8
|
||||
y = y // 8
|
||||
|
||||
#enfonce minimum size of 64
|
||||
if x > (samples.shape[3] - 8):
|
||||
x = samples.shape[3] - 8
|
||||
if y > (samples.shape[2] - 8):
|
||||
y = samples.shape[2] - 8
|
||||
|
||||
new_height = height // 8
|
||||
new_width = width // 8
|
||||
to_x = new_width + x
|
||||
to_y = new_height + y
|
||||
s['samples'] = samples[:,:,y:to_y, x:to_x]
|
||||
return s
|
||||
|
||||
def cleanup():
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
def cv2pil(cv_image):
|
||||
"""
|
||||
将OpenCV图像转换为PIL图像
|
||||
:param cv_image: OpenCV图像
|
||||
:return: PIL图像
|
||||
"""
|
||||
# 将图像从BGR转换为RGB
|
||||
rgb_image = cv2.cvtColor(cv_image, cv2.COLOR_BGR2RGB)
|
||||
# 使用PIL的Image.fromarray方法将NumPy数组转换为PIL图像
|
||||
pil_image = Image.fromarray(rgb_image)
|
||||
return pil_image
|
||||
|
||||
|
||||
def tensor_to_pil(tensor):
|
||||
image_np = tensor.squeeze().mul(255).clamp(0, 255).byte().numpy()
|
||||
image = Image.fromarray(image_np, mode='RGB')
|
||||
return image
|
||||
|
||||
def tensor2pil_list(image,width,height):
|
||||
B,_,_,_=image.size()
|
||||
if B==1:
|
||||
ref_image_list=[tensor2pil_upscale(image,width,height)]
|
||||
else:
|
||||
img_list = list(torch.chunk(image, chunks=B))
|
||||
ref_image_list = [tensor2pil_upscale(img,width,height) for img in img_list]
|
||||
return ref_image_list
|
||||
|
||||
|
||||
def tensor_upscale(img_tensor, width, height):
|
||||
samples = img_tensor.movedim(-1, 1)
|
||||
img = common_upscale(samples, width, height, "nearest-exact", "center")
|
||||
samples = img.movedim(1, -1)
|
||||
return samples
|
||||
|
||||
def tensor2pil_upscale(img_tensor, width, height):
|
||||
samples = img_tensor.movedim(-1, 1)
|
||||
img = common_upscale(samples, width, height, "nearest-exact", "center")
|
||||
samples = img.movedim(1, -1)
|
||||
img_pil = tensor_to_pil(samples)
|
||||
return img_pil
|
||||
|
||||
def tensor2np_upscale(img_tensor, width, height):
|
||||
samples = img_tensor.movedim(-1, 1)
|
||||
img = common_upscale(samples, width, height, "nearest-exact", "center")
|
||||
samples = img.movedim(1, -1)
|
||||
image_np = samples.squeeze().mul(255).clamp(0, 255).byte().numpy()
|
||||
return image_np
|
||||
|
||||
|
||||
def tensor2cv(tensor_image,RGB2BGR=True):
|
||||
if len(tensor_image.shape)==4:#bhwc to hwc
|
||||
tensor_image=tensor_image.squeeze(0)
|
||||
if tensor_image.is_cuda:
|
||||
tensor_image = tensor_image.cpu().detach()
|
||||
tensor_image=tensor_image.numpy()
|
||||
#反归一化
|
||||
maxValue=tensor_image.max()
|
||||
tensor_image=tensor_image*255/maxValue
|
||||
img_cv2=np.uint8(tensor_image)#32 to uint8
|
||||
if RGB2BGR:
|
||||
img_cv2=cv2.cvtColor(img_cv2,cv2.COLOR_RGB2BGR)
|
||||
return img_cv2
|
||||
|
||||
def cvargb2tensor(img):
|
||||
assert type(img) == np.ndarray, 'the img type is {}, but ndarry expected'.format(type(img))
|
||||
img = torch.from_numpy(img.transpose((2, 0, 1)))
|
||||
return img.float().div(255).unsqueeze(0) # 255也可以改为256
|
||||
|
||||
def cv2tensor(img):
|
||||
assert type(img) == np.ndarray, 'the img type is {}, but ndarry expected'.format(type(img))
|
||||
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
||||
img = torch.from_numpy(img.transpose((2, 0, 1)))
|
||||
return img.float().div(255).unsqueeze(0) # 255也可以改为256
|
||||
|
||||
def images_generator(img_list: list,):
|
||||
#get img size
|
||||
sizes = {}
|
||||
for image_ in img_list:
|
||||
if isinstance(image_,Image.Image):
|
||||
count = sizes.get(image_.size, 0)
|
||||
sizes[image_.size] = count + 1
|
||||
elif isinstance(image_,np.ndarray):
|
||||
count = sizes.get(image_.shape[:2][::-1], 0)
|
||||
sizes[image_.shape[:2][::-1]] = count + 1
|
||||
else:
|
||||
raise "unsupport image list,must be pil or cv2!!!"
|
||||
size = max(sizes.items(), key=lambda x: x[1])[0]
|
||||
yield size[0], size[1]
|
||||
|
||||
# any to tensor
|
||||
def load_image(img_in):
|
||||
if isinstance(img_in, Image.Image):
|
||||
img_in=img_in.convert("RGB")
|
||||
i = np.array(img_in, dtype=np.float32)
|
||||
i = torch.from_numpy(i).div_(255)
|
||||
if i.shape[0] != size[1] or i.shape[1] != size[0]:
|
||||
i = torch.from_numpy(i).movedim(-1, 0).unsqueeze(0)
|
||||
i = common_upscale(i, size[0], size[1], "lanczos", "center")
|
||||
i = i.squeeze(0).movedim(0, -1).numpy()
|
||||
return i
|
||||
elif isinstance(img_in,np.ndarray):
|
||||
i=cv2.cvtColor(img_in,cv2.COLOR_BGR2RGB).astype(np.float32)
|
||||
i = torch.from_numpy(i).div_(255)
|
||||
#print(i.shape)
|
||||
return i
|
||||
else:
|
||||
raise "unsupport image list,must be pil,cv2 or tensor!!!"
|
||||
|
||||
total_images = len(img_list)
|
||||
processed_images = 0
|
||||
pbar = ProgressBar(total_images)
|
||||
images = map(load_image, img_list)
|
||||
try:
|
||||
prev_image = next(images)
|
||||
while True:
|
||||
next_image = next(images)
|
||||
yield prev_image
|
||||
processed_images += 1
|
||||
pbar.update_absolute(processed_images, total_images)
|
||||
prev_image = next_image
|
||||
except StopIteration:
|
||||
pass
|
||||
if prev_image is not None:
|
||||
yield prev_image
|
||||
|
||||
def load_images(img_list: list,):
|
||||
gen = images_generator(img_list)
|
||||
(width, height) = next(gen)
|
||||
images = torch.from_numpy(np.fromiter(gen, np.dtype((np.float32, (height, width, 3)))))
|
||||
if len(images) == 0:
|
||||
raise FileNotFoundError(f"No images could be loaded .")
|
||||
return images
|
||||
|
||||
def tensor2pil(tensor):
|
||||
image_np = tensor.squeeze().mul(255).clamp(0, 255).byte().numpy()
|
||||
image = Image.fromarray(image_np, mode='RGB')
|
||||
return image
|
||||
|
||||
def pil2narry(img):
|
||||
narry = torch.from_numpy(np.array(img).astype(np.float32) / 255.0).unsqueeze(0)
|
||||
return narry
|
||||
|
||||
def equalize_lists(list1, list2):
|
||||
"""
|
||||
比较两个列表的长度，如果不一致，则将较短的列表复制以匹配较长列表的长度。
|
||||
|
||||
参数:
|
||||
list1 (list): 第一个列表
|
||||
list2 (list): 第二个列表
|
||||
|
||||
返回:
|
||||
tuple: 包含两个长度相等的列表的元组
|
||||
"""
|
||||
len1 = len(list1)
|
||||
len2 = len(list2)
|
||||
|
||||
if len1 == len2:
|
||||
pass
|
||||
elif len1 < len2:
|
||||
print("list1 is shorter than list2, copying list1 to match list2's length.")
|
||||
list1.extend(list1 * ((len2 // len1) + 1)) # 复制list1以匹配list2的长度
|
||||
list1 = list1[:len2] # 确保长度一致
|
||||
else:
|
||||
print("list2 is shorter than list1, copying list2 to match list1's length.")
|
||||
list2.extend(list2 * ((len1 // len2) + 1)) # 复制list2以匹配list1的长度
|
||||
list2 = list2[:len1] # 确保长度一致
|
||||
|
||||
return list1, list2
|
||||
|
||||
def file_exists(directory, filename):
|
||||
# 构建文件的完整路径
|
||||
file_path = os.path.join(directory, filename)
|
||||
# 检查文件是否存在
|
||||
return os.path.isfile(file_path)
|
||||
|
||||
def download_weights(file_dir,repo_id,subfolder="",pt_name=""):
|
||||
if subfolder:
|
||||
file_path = os.path.join(file_dir,subfolder, pt_name)
|
||||
sub_dir=os.path.join(file_dir,subfolder)
|
||||
if not os.path.exists(sub_dir):
|
||||
os.makedirs(sub_dir)
|
||||
if not os.path.exists(file_path):
|
||||
file_path = hf_hub_download(
|
||||
repo_id=repo_id,
|
||||
subfolder=subfolder,
|
||||
filename=pt_name,
|
||||
local_dir = file_dir,
|
||||
)
|
||||
return file_path
|
||||
else:
|
||||
file_path = os.path.join(file_dir, pt_name)
|
||||
if not os.path.exists(file_dir):
|
||||
os.makedirs(file_dir)
|
||||
if not os.path.exists(file_path):
|
||||
file_path = hf_hub_download(
|
||||
repo_id=repo_id,
|
||||
filename=pt_name,
|
||||
local_dir=file_dir,
|
||||
)
|
||||
return file_path
|
||||
+3
-3
@@ -1,9 +1,9 @@
|
||||
[project]
|
||||
name = "kv_edit"
|
||||
description = "KV-Edit: Training-Free Image Editing for Precise Background Preservation,you can use it in comfyUI"
|
||||
version = "1.0.0"
|
||||
description = "KV-Edit: Training-Free Image Editing for Precise Background Preservation - native ComfyUI FLUX DiT rewrite (no diffusers / transformers / accelerate)"
|
||||
version = "2.0.0"
|
||||
license = {file = "LICENSE"}
|
||||
dependencies = ["torch", "einops", "#fire", "#gradio==5.17.1", "#huggingface-hub", "#invisible-watermark", "#matplotlib", "numpy", "opencv-python", "Pillow", "Requests", "safetensors", "scikit-learn", "scipy", "scikit-image", "torchvision", "tqdm", "#transformers==4.49.0", "sentencepiece"]
|
||||
dependencies = ["torch"]
|
||||
|
||||
[project.urls]
|
||||
Repository = "https://github.com/smthemex/ComfyUI_KV_Edit"
|
||||
|
||||
@@ -1,19 +1 @@
|
||||
torch
|
||||
einops
|
||||
#fire
|
||||
#gradio==5.17.1
|
||||
#huggingface-hub
|
||||
#invisible-watermark
|
||||
#matplotlib
|
||||
numpy
|
||||
opencv-python
|
||||
Pillow
|
||||
Requests
|
||||
safetensors
|
||||
scikit-learn
|
||||
scipy
|
||||
scikit-image
|
||||
torchvision
|
||||
tqdm
|
||||
#transformers==4.49.0
|
||||
sentencepiece
|
||||
Binary file not shown.
|
Before Width: | Height: | Size: 494 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 428 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 424 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 727 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 4.4 MiB |
Reference in New Issue
Block a user