mirror of
https://github.com/lllyasviel/stable-diffusion-webui-forge.git
synced 2026-07-21 21:01:24 +08:00
Z-Image-Turbo
This commit is contained in:
Haoming
GitHub
parent
6a4dee6794
commit
93fd1ab26b
@ -130,7 +130,6 @@ dynamic_args = dict(
|
||||
kontext=False,
|
||||
edit=False,
|
||||
nunchaku=False,
|
||||
num_tokens={"c": None, "uc": None},
|
||||
ref_latents=[],
|
||||
concat_latent=None,
|
||||
)
|
||||
@ -141,7 +140,6 @@ Some parameters that are used throughout the Webui
|
||||
- kontext: `bool` - Flux Kontext
|
||||
- edit: `bool` - Qwen-Image-Edit
|
||||
- nunchaku: `bool` - Nunchaku (SVDQ) Models
|
||||
- num_tokens: `dict[str, int]` - Tokens for Lumina 2
|
||||
- ref_latents: `list[torch.Tensor]` - Reference Latent(s) for Flux Kontext & Qwen-Image-Edit
|
||||
- concat_latent: `torch.Tensor` - Input Latent for Wan 2.2 I2V
|
||||
"""
|
||||
|
||||
62
backend/diffusion_engine/zimage.py
Normal file
62
backend/diffusion_engine/zimage.py
Normal file
@ -0,0 +1,62 @@
|
||||
import torch
|
||||
from huggingface_guess import model_list
|
||||
|
||||
from backend import memory_management
|
||||
from backend.diffusion_engine.base import ForgeDiffusionEngine, ForgeObjects
|
||||
from backend.modules.k_prediction import PredictionDiscreteFlow
|
||||
from backend.patcher.clip import CLIP
|
||||
from backend.patcher.unet import UnetPatcher
|
||||
from backend.patcher.vae import VAE
|
||||
from backend.text_processing.qwen3_engine import Qwen3TextProcessingEngine
|
||||
|
||||
|
||||
class ZImage(ForgeDiffusionEngine):
|
||||
matched_guesses = [model_list.ZImage]
|
||||
|
||||
def __init__(self, estimated_config, huggingface_components):
|
||||
super().__init__(estimated_config, huggingface_components)
|
||||
self.is_inpaint = False
|
||||
|
||||
clip = CLIP(model_dict={"qwen3": huggingface_components["text_encoder"]}, tokenizer_dict={"qwen3": huggingface_components["tokenizer"]})
|
||||
|
||||
vae = VAE(model=huggingface_components["vae"])
|
||||
|
||||
k_predictor = PredictionDiscreteFlow(estimated_config)
|
||||
|
||||
unet = UnetPatcher.from_model(model=huggingface_components["transformer"], diffusers_scheduler=None, k_predictor=k_predictor, config=estimated_config)
|
||||
|
||||
self.text_processing_engine_gemma = Qwen3TextProcessingEngine(
|
||||
text_encoder=clip.cond_stage_model.qwen3,
|
||||
tokenizer=clip.tokenizer.qwen3,
|
||||
)
|
||||
|
||||
self.forge_objects = ForgeObjects(unet=unet, clip=clip, vae=vae, clipvision=None)
|
||||
self.forge_objects_original = self.forge_objects.shallow_copy()
|
||||
self.forge_objects_after_applying_lora = self.forge_objects.shallow_copy()
|
||||
|
||||
self.use_shift = True
|
||||
self.is_flux = True
|
||||
|
||||
@torch.inference_mode()
|
||||
def get_learned_conditioning(self, prompt: list[str]):
|
||||
memory_management.load_model_gpu(self.forge_objects.clip.patcher)
|
||||
shift = getattr(prompt, "distilled_cfg_scale", 3.0)
|
||||
self.forge_objects.unet.model.predictor.set_parameters(shift=shift)
|
||||
return self.text_processing_engine_gemma(prompt)
|
||||
|
||||
@torch.inference_mode()
|
||||
def get_prompt_lengths_on_ui(self, prompt):
|
||||
token_count = len(self.text_processing_engine_gemma.tokenize([prompt])[0])
|
||||
return token_count, max(999, token_count)
|
||||
|
||||
@torch.inference_mode()
|
||||
def encode_first_stage(self, x):
|
||||
sample = self.forge_objects.vae.encode(x.movedim(1, -1) * 0.5 + 0.5)
|
||||
sample = self.forge_objects.vae.first_stage_model.process_in(sample)
|
||||
return sample.to(x)
|
||||
|
||||
@torch.inference_mode()
|
||||
def decode_first_stage(self, x):
|
||||
sample = self.forge_objects.vae.first_stage_model.process_out(x)
|
||||
sample = self.forge_objects.vae.decode(sample).movedim(-1, 1) * 2.0 - 1.0
|
||||
return sample.to(x)
|
||||
@ -0,0 +1,24 @@
|
||||
{
|
||||
"_class_name": "ZImagePipeline",
|
||||
"_diffusers_version": "0.36.0.dev0",
|
||||
"scheduler": [
|
||||
"diffusers",
|
||||
"FlowMatchEulerDiscreteScheduler"
|
||||
],
|
||||
"text_encoder": [
|
||||
"transformers",
|
||||
"Qwen3Model"
|
||||
],
|
||||
"tokenizer": [
|
||||
"transformers",
|
||||
"Qwen2Tokenizer"
|
||||
],
|
||||
"transformer": [
|
||||
"diffusers",
|
||||
"ZImageTransformer2DModel"
|
||||
],
|
||||
"vae": [
|
||||
"diffusers",
|
||||
"AutoencoderKL"
|
||||
]
|
||||
}
|
||||
@ -0,0 +1,7 @@
|
||||
{
|
||||
"_class_name": "FlowMatchEulerDiscreteScheduler",
|
||||
"_diffusers_version": "0.36.0.dev0",
|
||||
"num_train_timesteps": 1000,
|
||||
"use_dynamic_shifting": false,
|
||||
"shift": 3.0
|
||||
}
|
||||
@ -0,0 +1,30 @@
|
||||
{
|
||||
"architectures": [
|
||||
"Qwen3ForCausalLM"
|
||||
],
|
||||
"attention_bias": false,
|
||||
"attention_dropout": 0.0,
|
||||
"bos_token_id": 151643,
|
||||
"eos_token_id": 151645,
|
||||
"head_dim": 128,
|
||||
"hidden_act": "silu",
|
||||
"hidden_size": 2560,
|
||||
"initializer_range": 0.02,
|
||||
"intermediate_size": 9728,
|
||||
"max_position_embeddings": 40960,
|
||||
"max_window_layers": 36,
|
||||
"model_type": "qwen3",
|
||||
"num_attention_heads": 32,
|
||||
"num_hidden_layers": 36,
|
||||
"num_key_value_heads": 8,
|
||||
"rms_norm_eps": 1e-06,
|
||||
"rope_scaling": null,
|
||||
"rope_theta": 1000000,
|
||||
"sliding_window": null,
|
||||
"tie_word_embeddings": true,
|
||||
"torch_dtype": "bfloat16",
|
||||
"transformers_version": "4.51.0",
|
||||
"use_cache": true,
|
||||
"use_sliding_window": false,
|
||||
"vocab_size": 151936
|
||||
}
|
||||
@ -0,0 +1,13 @@
|
||||
{
|
||||
"bos_token_id": 151643,
|
||||
"do_sample": true,
|
||||
"eos_token_id": [
|
||||
151645,
|
||||
151643
|
||||
],
|
||||
"pad_token_id": 151643,
|
||||
"temperature": 0.6,
|
||||
"top_k": 20,
|
||||
"top_p": 0.95,
|
||||
"transformers_version": "4.51.0"
|
||||
}
|
||||
151388
backend/huggingface/Tongyi-MAI/Z-Image-Turbo/tokenizer/merges.txt
Normal file
151388
backend/huggingface/Tongyi-MAI/Z-Image-Turbo/tokenizer/merges.txt
Normal file
File diff suppressed because it is too large
Load Diff
@ -0,0 +1,239 @@
|
||||
{
|
||||
"add_bos_token": false,
|
||||
"add_prefix_space": false,
|
||||
"added_tokens_decoder": {
|
||||
"151643": {
|
||||
"content": "<|endoftext|>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false,
|
||||
"special": true
|
||||
},
|
||||
"151644": {
|
||||
"content": "<|im_start|>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false,
|
||||
"special": true
|
||||
},
|
||||
"151645": {
|
||||
"content": "<|im_end|>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false,
|
||||
"special": true
|
||||
},
|
||||
"151646": {
|
||||
"content": "<|object_ref_start|>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false,
|
||||
"special": true
|
||||
},
|
||||
"151647": {
|
||||
"content": "<|object_ref_end|>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false,
|
||||
"special": true
|
||||
},
|
||||
"151648": {
|
||||
"content": "<|box_start|>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false,
|
||||
"special": true
|
||||
},
|
||||
"151649": {
|
||||
"content": "<|box_end|>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false,
|
||||
"special": true
|
||||
},
|
||||
"151650": {
|
||||
"content": "<|quad_start|>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false,
|
||||
"special": true
|
||||
},
|
||||
"151651": {
|
||||
"content": "<|quad_end|>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false,
|
||||
"special": true
|
||||
},
|
||||
"151652": {
|
||||
"content": "<|vision_start|>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false,
|
||||
"special": true
|
||||
},
|
||||
"151653": {
|
||||
"content": "<|vision_end|>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false,
|
||||
"special": true
|
||||
},
|
||||
"151654": {
|
||||
"content": "<|vision_pad|>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false,
|
||||
"special": true
|
||||
},
|
||||
"151655": {
|
||||
"content": "<|image_pad|>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false,
|
||||
"special": true
|
||||
},
|
||||
"151656": {
|
||||
"content": "<|video_pad|>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false,
|
||||
"special": true
|
||||
},
|
||||
"151657": {
|
||||
"content": "<tool_call>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false,
|
||||
"special": false
|
||||
},
|
||||
"151658": {
|
||||
"content": "</tool_call>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false,
|
||||
"special": false
|
||||
},
|
||||
"151659": {
|
||||
"content": "<|fim_prefix|>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false,
|
||||
"special": false
|
||||
},
|
||||
"151660": {
|
||||
"content": "<|fim_middle|>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false,
|
||||
"special": false
|
||||
},
|
||||
"151661": {
|
||||
"content": "<|fim_suffix|>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false,
|
||||
"special": false
|
||||
},
|
||||
"151662": {
|
||||
"content": "<|fim_pad|>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false,
|
||||
"special": false
|
||||
},
|
||||
"151663": {
|
||||
"content": "<|repo_name|>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false,
|
||||
"special": false
|
||||
},
|
||||
"151664": {
|
||||
"content": "<|file_sep|>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false,
|
||||
"special": false
|
||||
},
|
||||
"151665": {
|
||||
"content": "<tool_response>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false,
|
||||
"special": false
|
||||
},
|
||||
"151666": {
|
||||
"content": "</tool_response>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false,
|
||||
"special": false
|
||||
},
|
||||
"151667": {
|
||||
"content": "<think>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false,
|
||||
"special": false
|
||||
},
|
||||
"151668": {
|
||||
"content": "</think>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false,
|
||||
"special": false
|
||||
}
|
||||
},
|
||||
"additional_special_tokens": [
|
||||
"<|im_start|>",
|
||||
"<|im_end|>",
|
||||
"<|object_ref_start|>",
|
||||
"<|object_ref_end|>",
|
||||
"<|box_start|>",
|
||||
"<|box_end|>",
|
||||
"<|quad_start|>",
|
||||
"<|quad_end|>",
|
||||
"<|vision_start|>",
|
||||
"<|vision_end|>",
|
||||
"<|vision_pad|>",
|
||||
"<|image_pad|>",
|
||||
"<|video_pad|>"
|
||||
],
|
||||
"bos_token": null,
|
||||
"chat_template": "{%- if tools %}\n {{- '<|im_start|>system\\n' }}\n {%- if messages[0].role == 'system' %}\n {{- messages[0].content + '\\n\\n' }}\n {%- endif %}\n {{- \"# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>\" }}\n {%- for tool in tools %}\n {{- \"\\n\" }}\n {{- tool | tojson }}\n {%- endfor %}\n {{- \"\\n</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\\\"name\\\": <function-name>, \\\"arguments\\\": <args-json-object>}\\n</tool_call><|im_end|>\\n\" }}\n{%- else %}\n {%- if messages[0].role == 'system' %}\n {{- '<|im_start|>system\\n' + messages[0].content + '<|im_end|>\\n' }}\n {%- endif %}\n{%- endif %}\n{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}\n{%- for message in messages[::-1] %}\n {%- set index = (messages|length - 1) - loop.index0 %}\n {%- if ns.multi_step_tool and message.role == \"user\" and message.content is string and not(message.content.startswith('<tool_response>') and message.content.endswith('</tool_response>')) %}\n {%- set ns.multi_step_tool = false %}\n {%- set ns.last_query_index = index %}\n {%- endif %}\n{%- endfor %}\n{%- for message in messages %}\n {%- if message.content is string %}\n {%- set content = message.content %}\n {%- else %}\n {%- set content = '' %}\n {%- endif %}\n {%- if (message.role == \"user\") or (message.role == \"system\" and not loop.first) %}\n {{- '<|im_start|>' + message.role + '\\n' + content + '<|im_end|>' + '\\n' }}\n {%- elif message.role == \"assistant\" %}\n {%- set reasoning_content = '' %}\n {%- if message.reasoning_content is string %}\n {%- set reasoning_content = message.reasoning_content %}\n {%- else %}\n {%- if '</think>' in content %}\n {%- set reasoning_content = content.split('</think>')[0].rstrip('\\n').split('<think>')[-1].lstrip('\\n') %}\n {%- set content = content.split('</think>')[-1].lstrip('\\n') %}\n {%- endif %}\n {%- endif %}\n {%- if loop.index0 > ns.last_query_index %}\n {%- if loop.last or (not loop.last and reasoning_content) %}\n {{- '<|im_start|>' + message.role + '\\n<think>\\n' + reasoning_content.strip('\\n') + '\\n</think>\\n\\n' + content.lstrip('\\n') }}\n {%- else %}\n {{- '<|im_start|>' + message.role + '\\n' + content }}\n {%- endif %}\n {%- else %}\n {{- '<|im_start|>' + message.role + '\\n' + content }}\n {%- endif %}\n {%- if message.tool_calls %}\n {%- for tool_call in message.tool_calls %}\n {%- if (loop.first and content) or (not loop.first) %}\n {{- '\\n' }}\n {%- endif %}\n {%- if tool_call.function %}\n {%- set tool_call = tool_call.function %}\n {%- endif %}\n {{- '<tool_call>\\n{\"name\": \"' }}\n {{- tool_call.name }}\n {{- '\", \"arguments\": ' }}\n {%- if tool_call.arguments is string %}\n {{- tool_call.arguments }}\n {%- else %}\n {{- tool_call.arguments | tojson }}\n {%- endif %}\n {{- '}\\n</tool_call>' }}\n {%- endfor %}\n {%- endif %}\n {{- '<|im_end|>\\n' }}\n {%- elif message.role == \"tool\" %}\n {%- if loop.first or (messages[loop.index0 - 1].role != \"tool\") %}\n {{- '<|im_start|>user' }}\n {%- endif %}\n {{- '\\n<tool_response>\\n' }}\n {{- content }}\n {{- '\\n</tool_response>' }}\n {%- if loop.last or (messages[loop.index0 + 1].role != \"tool\") %}\n {{- '<|im_end|>\\n' }}\n {%- endif %}\n {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n {{- '<|im_start|>assistant\\n' }}\n {%- if enable_thinking is defined and enable_thinking is false %}\n {{- '<think>\\n\\n</think>\\n\\n' }}\n {%- endif %}\n{%- endif %}",
|
||||
"clean_up_tokenization_spaces": false,
|
||||
"eos_token": "<|im_end|>",
|
||||
"errors": "replace",
|
||||
"model_max_length": 131072,
|
||||
"pad_token": "<|endoftext|>",
|
||||
"split_special_tokens": false,
|
||||
"tokenizer_class": "Qwen2Tokenizer",
|
||||
"unk_token": null
|
||||
}
|
||||
File diff suppressed because one or more lines are too long
@ -0,0 +1,31 @@
|
||||
{
|
||||
"_class_name": "ZImageTransformer2DModel",
|
||||
"_diffusers_version": "0.36.0.dev0",
|
||||
"all_f_patch_size": [
|
||||
1
|
||||
],
|
||||
"all_patch_size": [
|
||||
2
|
||||
],
|
||||
"axes_dims": [
|
||||
32,
|
||||
48,
|
||||
48
|
||||
],
|
||||
"axes_lens": [
|
||||
1536,
|
||||
512,
|
||||
512
|
||||
],
|
||||
"cap_feat_dim": 2560,
|
||||
"dim": 3840,
|
||||
"in_channels": 16,
|
||||
"n_heads": 30,
|
||||
"n_kv_heads": 30,
|
||||
"n_layers": 30,
|
||||
"n_refiner_layers": 2,
|
||||
"norm_eps": 1e-05,
|
||||
"qk_norm": true,
|
||||
"rope_theta": 256.0,
|
||||
"t_scale": 1000.0
|
||||
}
|
||||
38
backend/huggingface/Tongyi-MAI/Z-Image-Turbo/vae/config.json
Normal file
38
backend/huggingface/Tongyi-MAI/Z-Image-Turbo/vae/config.json
Normal file
@ -0,0 +1,38 @@
|
||||
{
|
||||
"_class_name": "AutoencoderKL",
|
||||
"_diffusers_version": "0.36.0.dev0",
|
||||
"_name_or_path": "flux-dev",
|
||||
"act_fn": "silu",
|
||||
"block_out_channels": [
|
||||
128,
|
||||
256,
|
||||
512,
|
||||
512
|
||||
],
|
||||
"down_block_types": [
|
||||
"DownEncoderBlock2D",
|
||||
"DownEncoderBlock2D",
|
||||
"DownEncoderBlock2D",
|
||||
"DownEncoderBlock2D"
|
||||
],
|
||||
"force_upcast": true,
|
||||
"in_channels": 3,
|
||||
"latent_channels": 16,
|
||||
"latents_mean": null,
|
||||
"latents_std": null,
|
||||
"layers_per_block": 2,
|
||||
"mid_block_add_attention": true,
|
||||
"norm_num_groups": 32,
|
||||
"out_channels": 3,
|
||||
"sample_size": 1024,
|
||||
"scaling_factor": 0.3611,
|
||||
"shift_factor": 0.1159,
|
||||
"up_block_types": [
|
||||
"UpDecoderBlock2D",
|
||||
"UpDecoderBlock2D",
|
||||
"UpDecoderBlock2D",
|
||||
"UpDecoderBlock2D"
|
||||
],
|
||||
"use_post_quant_conv": false,
|
||||
"use_quant_conv": false
|
||||
}
|
||||
@ -11,6 +11,9 @@ class Token:
|
||||
neta_compress = os.path.join(folder, "neta.tokenizer.json.xz")
|
||||
neta_lumina = os.path.join(folder, "neta-art", "Neta-Lumina", "tokenizer", "tokenizer.json")
|
||||
|
||||
z_compress = os.path.join(folder, "z.tokenizer.json.xz")
|
||||
z_image = os.path.join(folder, "Tongyi-MAI", "Z-Image-Turbo", "tokenizer", "tokenizer.json")
|
||||
|
||||
|
||||
class sha256:
|
||||
wan = "20a46ac256746594ed7e1e3ef733b83fbc5a6f0922aa7480eda961743de080ef"
|
||||
@ -19,6 +22,9 @@ class sha256:
|
||||
neta = "3f289bc05132635a8bc7aca7aa21255efd5e18f3710f43e3cdb96bcd41be4922"
|
||||
# https://huggingface.co/neta-art/Neta-Lumina-diffusers/blob/main/tokenizer/tokenizer.json
|
||||
|
||||
z = "aeb13307a71acd8fe81861d94ad54ab689df773318809eed3cbe794b4492dae4"
|
||||
# https://huggingface.co/Tongyi-MAI/Z-Image-Turbo/blob/main/tokenizer/tokenizer.json
|
||||
|
||||
|
||||
def decompress(source: str, target: str):
|
||||
import lzma
|
||||
@ -46,12 +52,17 @@ def process():
|
||||
if not os.path.isfile(Token.neta_lumina):
|
||||
decompress(Token.neta_compress, Token.neta_lumina)
|
||||
compare_sha256(Token.neta_lumina, sha256.neta)
|
||||
if not os.path.isfile(Token.z_image):
|
||||
decompress(Token.z_compress, Token.z_image)
|
||||
compare_sha256(Token.z_image, sha256.z)
|
||||
|
||||
# if not os.path.isfile(Token.wan_compress):
|
||||
# compress(Token.wan_t2v, Token.wan_compress)
|
||||
# if not os.path.isfile(Token.neta_compress):
|
||||
# compare_sha256(Token.neta_lumina, sha256.neta)
|
||||
# compress(Token.neta_lumina, Token.neta_compress)
|
||||
# if not os.path.isfile(Token.z_compress):
|
||||
# compress(Token.z_image, Token.z_compress)
|
||||
|
||||
|
||||
def compare_sha256(path: str, target: str):
|
||||
|
||||
BIN
backend/huggingface/z.tokenizer.json.xz
Normal file
BIN
backend/huggingface/z.tokenizer.json.xz
Normal file
Binary file not shown.
@ -16,6 +16,7 @@ from backend.diffusion_engine.qwen import QwenImage
|
||||
from backend.diffusion_engine.sd15 import StableDiffusion
|
||||
from backend.diffusion_engine.sdxl import StableDiffusionXL, StableDiffusionXLRefiner
|
||||
from backend.diffusion_engine.wan import Wan
|
||||
from backend.diffusion_engine.zimage import ZImage
|
||||
from backend.nn.clip import IntegratedCLIP
|
||||
from backend.nn.unet import IntegratedUNet2DConditionModel
|
||||
from backend.nn.vae import IntegratedAutoencoderKL
|
||||
@ -28,7 +29,7 @@ from backend.utils import (
|
||||
read_arbitrary_config,
|
||||
)
|
||||
|
||||
possible_models = [StableDiffusion, StableDiffusionXLRefiner, StableDiffusionXL, Chroma, Flux, Wan, QwenImage, Lumina2]
|
||||
possible_models = [StableDiffusion, StableDiffusionXLRefiner, StableDiffusionXL, Chroma, Flux, Wan, QwenImage, Lumina2, ZImage]
|
||||
|
||||
|
||||
logging.getLogger("diffusers").setLevel(logging.ERROR)
|
||||
@ -151,6 +152,38 @@ def load_huggingface_component(guess, component_name, lib_name, cls_name, repo_p
|
||||
|
||||
load_state_dict(model, state_dict, log_name=cls_name, ignore_errors=[])
|
||||
|
||||
return model
|
||||
if cls_name == "Qwen3Model":
|
||||
assert isinstance(state_dict, dict) and len(state_dict) > 16, "You do not have Qwen3 state dict!"
|
||||
|
||||
from backend.nn.llm.llama import Qwen3_4B
|
||||
|
||||
config = read_arbitrary_config(config_path)
|
||||
|
||||
storage_dtype = memory_management.text_encoder_dtype()
|
||||
state_dict_dtype = memory_management.state_dict_dtype(state_dict)
|
||||
|
||||
if state_dict_dtype in [torch.float8_e4m3fn, torch.float8_e5m2, "nf4", "fp4", "gguf"]:
|
||||
print(f"Using Detected Qwen3 Data Type: {state_dict_dtype}")
|
||||
storage_dtype = state_dict_dtype
|
||||
if state_dict_dtype in ["nf4", "fp4", "gguf"]:
|
||||
print("Using pre-quant state dict!")
|
||||
if state_dict_dtype in ["gguf"]:
|
||||
beautiful_print_gguf_state_dict_statics(state_dict)
|
||||
else:
|
||||
print(f"Using Default Qwen3 Data Type: {storage_dtype}")
|
||||
|
||||
if storage_dtype in ["nf4", "fp4", "gguf"]:
|
||||
with modeling_utils.no_init_weights():
|
||||
with using_forge_operations(device=memory_management.cpu, dtype=memory_management.text_encoder_dtype(), manual_cast_enabled=False, bnb_dtype=storage_dtype):
|
||||
model = Qwen3_4B(config)
|
||||
else:
|
||||
with modeling_utils.no_init_weights():
|
||||
with using_forge_operations(device=memory_management.cpu, dtype=storage_dtype, manual_cast_enabled=True):
|
||||
model = Qwen3_4B(config)
|
||||
|
||||
load_state_dict(model, state_dict, log_name=cls_name, ignore_errors=[])
|
||||
|
||||
return model
|
||||
if cls_name in ["T5EncoderModel", "UMT5EncoderModel"]:
|
||||
assert isinstance(state_dict, dict) and len(state_dict) > 16, "You do not have T5 state dict!"
|
||||
@ -194,7 +227,7 @@ def load_huggingface_component(guess, component_name, lib_name, cls_name, repo_p
|
||||
load_state_dict(model, state_dict, log_name=cls_name, ignore_errors=["transformer.encoder.embed_tokens.weight", "logit_scale"])
|
||||
|
||||
return model
|
||||
if cls_name in ["UNet2DConditionModel", "FluxTransformer2DModel", "ChromaTransformer2DModel", "WanTransformer3DModel", "QwenImageTransformer2DModel", "Lumina2Transformer2DModel"]:
|
||||
if cls_name in ["UNet2DConditionModel", "FluxTransformer2DModel", "ChromaTransformer2DModel", "WanTransformer3DModel", "QwenImageTransformer2DModel", "Lumina2Transformer2DModel", "ZImageTransformer2DModel"]:
|
||||
assert isinstance(state_dict, dict) and len(state_dict) > 16, "You do not have model state dict!"
|
||||
|
||||
model_loader = None
|
||||
@ -226,7 +259,7 @@ def load_huggingface_component(guess, component_name, lib_name, cls_name, repo_p
|
||||
from backend.nn.qwen import QwenImageTransformer2DModel
|
||||
|
||||
model_loader = lambda c: QwenImageTransformer2DModel(**c)
|
||||
elif cls_name == "Lumina2Transformer2DModel":
|
||||
elif cls_name in ("Lumina2Transformer2DModel", "ZImageTransformer2DModel"):
|
||||
from backend.nn.lumina import NextDiT
|
||||
|
||||
model_loader = lambda c: NextDiT(**c)
|
||||
@ -529,12 +562,6 @@ def replace_state_dict(sd: dict[str, torch.Tensor], asd: dict[str, torch.Tensor]
|
||||
sd[f"{text_encoder_key_prefix}t5xxl.transformer.{k}"] = True
|
||||
sd[f"{text_encoder_key_prefix}t5xxl.transformer.filename"] = str(path)
|
||||
|
||||
if "model.layers.0.self_attn.k_proj.bias" in asd:
|
||||
weight = asd["model.layers.0.self_attn.k_proj.bias"]
|
||||
assert weight.shape[0] == 512
|
||||
for k, v in asd.items():
|
||||
sd[f"{text_encoder_key_prefix}qwen25_7b.{k}"] = v
|
||||
|
||||
if "model.layers.0.post_feedforward_layernorm.weight" in asd:
|
||||
assert "model.layers.0.self_attn.q_norm.weight" not in asd
|
||||
for k, v in asd.items():
|
||||
@ -542,6 +569,17 @@ def replace_state_dict(sd: dict[str, torch.Tensor], asd: dict[str, torch.Tensor]
|
||||
continue
|
||||
sd[f"{text_encoder_key_prefix}gemma2_2b.{k}"] = v
|
||||
|
||||
elif "model.layers.0.self_attn.k_proj.bias" in asd:
|
||||
weight = asd["model.layers.0.self_attn.k_proj.bias"]
|
||||
assert weight.shape[0] == 512
|
||||
for k, v in asd.items():
|
||||
sd[f"{text_encoder_key_prefix}qwen25_7b.{k}"] = v
|
||||
|
||||
elif "model.layers.0.post_attention_layernorm.weight" in asd:
|
||||
assert "model.layers.0.self_attn.q_norm.weight" in asd
|
||||
for k, v in asd.items():
|
||||
sd[f"{text_encoder_key_prefix}qwen3_4b.transformer.{k}"] = v
|
||||
|
||||
return sd
|
||||
|
||||
|
||||
@ -671,7 +709,7 @@ def forge_loader(sd: os.PathLike, additional_state_dicts: list[os.PathLike] = No
|
||||
huggingface_components["scheduler"].config.prediction_type = prediction_types.get(estimated_config.model_type.name, huggingface_components["scheduler"].config.prediction_type)
|
||||
|
||||
for M in possible_models:
|
||||
if any(isinstance(estimated_config, x) for x in M.matched_guesses):
|
||||
if any(type(estimated_config) is x for x in M.matched_guesses):
|
||||
return M(estimated_config=estimated_config, huggingface_components=huggingface_components)
|
||||
|
||||
print("Failed to recognize model type!")
|
||||
|
||||
@ -1,4 +1,4 @@
|
||||
# https://github.com/comfyanonymous/ComfyUI/blob/v0.3.64/comfy/text_encoders/llama.py
|
||||
# https://github.com/comfyanonymous/ComfyUI/blob/v0.3.75/comfy/text_encoders/llama.py
|
||||
|
||||
import math
|
||||
from dataclasses import asdict, dataclass
|
||||
@ -8,7 +8,6 @@ import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from backend.memory_management import (
|
||||
cast_to_device,
|
||||
is_device_cpu,
|
||||
text_encoder_device,
|
||||
xformers_enabled,
|
||||
@ -22,6 +21,29 @@ else:
|
||||
from . import qwen_vl
|
||||
|
||||
|
||||
@dataclass
|
||||
class Qwen3_4BConfig:
|
||||
vocab_size: int = 151936
|
||||
hidden_size: int = 2560
|
||||
intermediate_size: int = 9728
|
||||
num_hidden_layers: int = 36
|
||||
num_attention_heads: int = 32
|
||||
num_key_value_heads: int = 8
|
||||
max_position_embeddings: int = 40960
|
||||
rms_norm_eps: float = 1e-6
|
||||
rope_theta: float = 1000000.0
|
||||
transformer_type: str = "llama"
|
||||
head_dim = 128
|
||||
rms_norm_add = False
|
||||
mlp_activation = "silu"
|
||||
qkv_bias = False
|
||||
rope_dims = None
|
||||
q_norm = "gemma3"
|
||||
k_norm = "gemma3"
|
||||
rope_scale = None
|
||||
final_norm: bool = True
|
||||
|
||||
|
||||
@dataclass
|
||||
class Qwen25_7BVLI_Config:
|
||||
vocab_size: int = 152064
|
||||
@ -39,6 +61,10 @@ class Qwen25_7BVLI_Config:
|
||||
mlp_activation = "silu"
|
||||
qkv_bias = True
|
||||
rope_dims = [16, 24, 24]
|
||||
q_norm = None
|
||||
k_norm = None
|
||||
rope_scale = None
|
||||
final_norm: bool = True
|
||||
|
||||
|
||||
@dataclass
|
||||
@ -58,25 +84,11 @@ class Gemma2_2B_Config:
|
||||
mlp_activation = "gelu_pytorch_tanh"
|
||||
qkv_bias = False
|
||||
rope_dims = None
|
||||
|
||||
|
||||
def _rms_norm(x, weight, eps):
|
||||
return nn.functional.rms_norm(x, weight.shape, weight=cast_to_device(weight, device=x.device, dtype=x.dtype), eps=eps)
|
||||
|
||||
|
||||
class RMSNorm(nn.Module):
|
||||
def __init__(self, dim: int, eps: float, add=False):
|
||||
super().__init__()
|
||||
self.weight = nn.Parameter(torch.empty(dim))
|
||||
self.eps = eps
|
||||
self.add = add
|
||||
|
||||
def forward(self, x: torch.Tensor):
|
||||
w = self.weight
|
||||
if self.add:
|
||||
w = w + 1.0
|
||||
|
||||
return _rms_norm(x, w, self.eps)
|
||||
q_norm = None
|
||||
k_norm = None
|
||||
sliding_attention = None
|
||||
rope_scale = None
|
||||
final_norm: bool = True
|
||||
|
||||
|
||||
def rotate_half(x):
|
||||
@ -86,15 +98,21 @@ def rotate_half(x):
|
||||
return torch.cat((-x2, x1), dim=-1)
|
||||
|
||||
|
||||
def precompute_freqs_cis(head_dim, position_ids, theta, rope_dims=None, device=None):
|
||||
def precompute_freqs_cis(head_dim, position_ids, theta, rope_scale=None, rope_dims=None, device=None):
|
||||
if not isinstance(theta, list):
|
||||
theta = [theta]
|
||||
|
||||
out = []
|
||||
for t in theta:
|
||||
for index, t in enumerate(theta):
|
||||
theta_numerator = torch.arange(0, head_dim, 2, device=device).float()
|
||||
inv_freq = 1.0 / (t ** (theta_numerator / head_dim))
|
||||
|
||||
if rope_scale is not None:
|
||||
if isinstance(rope_scale, list):
|
||||
inv_freq /= rope_scale[index]
|
||||
else:
|
||||
inv_freq /= rope_scale
|
||||
|
||||
inv_freq_expanded = inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)
|
||||
position_ids_expanded = position_ids[:, None, :].float()
|
||||
freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
|
||||
@ -126,7 +144,7 @@ def apply_rope(xq, xk, freqs_cis):
|
||||
|
||||
|
||||
class Attention(nn.Module):
|
||||
def __init__(self, config: Qwen25_7BVLI_Config | Gemma2_2B_Config):
|
||||
def __init__(self, config: Qwen3_4BConfig | Qwen25_7BVLI_Config | Gemma2_2B_Config):
|
||||
super().__init__()
|
||||
self.num_heads = config.num_attention_heads
|
||||
self.num_kv_heads = config.num_key_value_heads
|
||||
@ -140,6 +158,16 @@ class Attention(nn.Module):
|
||||
self.v_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=config.qkv_bias)
|
||||
self.o_proj = nn.Linear(self.inner_size, config.hidden_size, bias=False)
|
||||
|
||||
if config.q_norm == "gemma3":
|
||||
self.q_norm = nn.RMSNorm(self.head_dim, eps=config.rms_norm_eps, add=config.rms_norm_add)
|
||||
else:
|
||||
self.q_norm = None
|
||||
|
||||
if config.k_norm == "gemma3":
|
||||
self.k_norm = nn.RMSNorm(self.head_dim, eps=config.rms_norm_eps, add=config.rms_norm_add)
|
||||
else:
|
||||
self.k_norm = None
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
@ -156,6 +184,11 @@ class Attention(nn.Module):
|
||||
xk = xk.view(batch_size, seq_length, self.num_kv_heads, self.head_dim).transpose(1, 2)
|
||||
xv = xv.view(batch_size, seq_length, self.num_kv_heads, self.head_dim).transpose(1, 2)
|
||||
|
||||
if self.q_norm is not None:
|
||||
xq = self.q_norm(xq)
|
||||
if self.k_norm is not None:
|
||||
xk = self.k_norm(xk)
|
||||
|
||||
xq, xk = apply_rope(xq, xk, freqs_cis=freqs_cis)
|
||||
|
||||
xk = xk.repeat_interleave(self.num_heads // self.num_kv_heads, dim=1)
|
||||
@ -166,7 +199,7 @@ class Attention(nn.Module):
|
||||
|
||||
|
||||
class MLP(nn.Module):
|
||||
def __init__(self, config: Qwen25_7BVLI_Config | Gemma2_2B_Config):
|
||||
def __init__(self, config: Qwen3_4BConfig | Qwen25_7BVLI_Config | Gemma2_2B_Config):
|
||||
super().__init__()
|
||||
self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
|
||||
self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
|
||||
@ -181,12 +214,12 @@ class MLP(nn.Module):
|
||||
|
||||
|
||||
class TransformerBlock(nn.Module):
|
||||
def __init__(self, config: Qwen25_7BVLI_Config | Gemma2_2B_Config, index):
|
||||
def __init__(self, config: Qwen3_4BConfig | Qwen25_7BVLI_Config | Gemma2_2B_Config, index):
|
||||
super().__init__()
|
||||
self.self_attn = Attention(config)
|
||||
self.mlp = MLP(config)
|
||||
self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
||||
self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
||||
self.input_layernorm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
||||
self.post_attention_layernorm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
@ -216,15 +249,20 @@ class TransformerBlock(nn.Module):
|
||||
|
||||
|
||||
class TransformerBlockGemma2(nn.Module):
|
||||
def __init__(self, config: Qwen25_7BVLI_Config | Gemma2_2B_Config, index):
|
||||
def __init__(self, config: Qwen3_4BConfig | Qwen25_7BVLI_Config | Gemma2_2B_Config, index):
|
||||
super().__init__()
|
||||
self.self_attn = Attention(config)
|
||||
self.mlp = MLP(config)
|
||||
self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps, add=config.rms_norm_add)
|
||||
self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps, add=config.rms_norm_add)
|
||||
self.pre_feedforward_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps, add=config.rms_norm_add)
|
||||
self.post_feedforward_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps, add=config.rms_norm_add)
|
||||
self.sliding_attention = False
|
||||
self.input_layernorm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps, add=config.rms_norm_add)
|
||||
self.post_attention_layernorm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps, add=config.rms_norm_add)
|
||||
self.pre_feedforward_layernorm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps, add=config.rms_norm_add)
|
||||
self.post_feedforward_layernorm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps, add=config.rms_norm_add)
|
||||
|
||||
if config.sliding_attention is not None:
|
||||
self.sliding_attention = config.sliding_attention[index % len(config.sliding_attention)]
|
||||
else:
|
||||
self.sliding_attention = False
|
||||
|
||||
self.transformer_type = config.transformer_type
|
||||
|
||||
def forward(
|
||||
@ -234,6 +272,13 @@ class TransformerBlockGemma2(nn.Module):
|
||||
freqs_cis: Optional[torch.Tensor] = None,
|
||||
optimized_attention=None,
|
||||
):
|
||||
if self.transformer_type == "gemma3":
|
||||
if self.sliding_attention:
|
||||
assert x.shape[1] <= self.sliding_attention
|
||||
freqs_cis = freqs_cis[1]
|
||||
else:
|
||||
freqs_cis = freqs_cis[0]
|
||||
|
||||
# Self Attention
|
||||
residual = x
|
||||
x = self.input_layernorm(x)
|
||||
@ -258,13 +303,13 @@ class TransformerBlockGemma2(nn.Module):
|
||||
|
||||
|
||||
class Llama2_(nn.Module):
|
||||
def __init__(self, config: Qwen25_7BVLI_Config | Gemma2_2B_Config):
|
||||
def __init__(self, config: Qwen3_4BConfig | Qwen25_7BVLI_Config | Gemma2_2B_Config):
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.vocab_size = config.vocab_size
|
||||
|
||||
self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
|
||||
if self.config.transformer_type == "gemma2":
|
||||
if self.config.transformer_type == "gemma2" or self.config.transformer_type == "gemma3":
|
||||
transformer = TransformerBlockGemma2
|
||||
self.normalize_in = True
|
||||
else:
|
||||
@ -272,13 +317,17 @@ class Llama2_(nn.Module):
|
||||
self.normalize_in = False
|
||||
|
||||
self.layers = nn.ModuleList([transformer(config, index=i) for i in range(config.num_hidden_layers)])
|
||||
self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps, add=config.rms_norm_add)
|
||||
|
||||
if config.final_norm:
|
||||
self.norm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps, add=config.rms_norm_add)
|
||||
else:
|
||||
self.norm = None
|
||||
|
||||
def forward(self, x, attention_mask=None, embeds=None, num_tokens=None, intermediate_output=None, final_layer_norm_intermediate=True, dtype=None, position_ids=None, embeds_info=[]):
|
||||
if embeds is not None:
|
||||
x = embeds
|
||||
else:
|
||||
x = self.embed_tokens(x)
|
||||
x = self.embed_tokens(x, out_dtype=dtype)
|
||||
|
||||
if self.normalize_in:
|
||||
x *= self.config.hidden_size**0.5
|
||||
@ -286,7 +335,7 @@ class Llama2_(nn.Module):
|
||||
if position_ids is None:
|
||||
position_ids = torch.arange(0, x.shape[1], device=x.device).unsqueeze(0)
|
||||
|
||||
freqs_cis = precompute_freqs_cis(self.config.head_dim, position_ids, self.config.rope_theta, self.config.rope_dims, device=x.device)
|
||||
freqs_cis = precompute_freqs_cis(self.config.head_dim, position_ids, self.config.rope_theta, self.config.rope_scale, self.config.rope_dims, device=x.device)
|
||||
|
||||
mask = None
|
||||
if attention_mask is not None:
|
||||
@ -301,8 +350,12 @@ class Llama2_(nn.Module):
|
||||
|
||||
intermediate = None
|
||||
all_intermediate = None
|
||||
only_layers = None
|
||||
if intermediate_output is not None:
|
||||
if intermediate_output == "all":
|
||||
if isinstance(intermediate_output, list):
|
||||
all_intermediate = []
|
||||
only_layers = set(intermediate_output)
|
||||
elif intermediate_output == "all":
|
||||
all_intermediate = []
|
||||
intermediate_output = None
|
||||
elif intermediate_output < 0:
|
||||
@ -310,7 +363,8 @@ class Llama2_(nn.Module):
|
||||
|
||||
for i, layer in enumerate(self.layers):
|
||||
if all_intermediate is not None:
|
||||
all_intermediate.append(x.unsqueeze(1).clone())
|
||||
if only_layers is None or (i in only_layers):
|
||||
all_intermediate.append(x.unsqueeze(1).clone())
|
||||
x = layer(
|
||||
x=x,
|
||||
attention_mask=mask,
|
||||
@ -320,14 +374,17 @@ class Llama2_(nn.Module):
|
||||
if i == intermediate_output:
|
||||
intermediate = x.clone()
|
||||
|
||||
x = self.norm(x)
|
||||
if self.norm is not None:
|
||||
x = self.norm(x)
|
||||
|
||||
if all_intermediate is not None:
|
||||
all_intermediate.append(x.unsqueeze(1).clone())
|
||||
if only_layers is None or ((i + 1) in only_layers):
|
||||
all_intermediate.append(x.unsqueeze(1).clone())
|
||||
|
||||
if all_intermediate is not None:
|
||||
intermediate = torch.cat(all_intermediate, dim=1)
|
||||
|
||||
if intermediate is not None and final_layer_norm_intermediate:
|
||||
if intermediate is not None and final_layer_norm_intermediate and self.norm is not None:
|
||||
intermediate = self.norm(intermediate)
|
||||
|
||||
return x, intermediate
|
||||
@ -344,6 +401,21 @@ class BaseLlama:
|
||||
return self.model(input_ids, *args, **kwargs)
|
||||
|
||||
|
||||
class Qwen3_4B(BaseLlama, nn.Module):
|
||||
def __init__(self, config_dict):
|
||||
super().__init__()
|
||||
config = Qwen3_4BConfig()
|
||||
|
||||
_config_dict = asdict(config)
|
||||
for key, value in _config_dict.items():
|
||||
if key in config_dict:
|
||||
assert value == config_dict[key]
|
||||
|
||||
self.num_layers = config.num_hidden_layers
|
||||
|
||||
self.model = Llama2_(config)
|
||||
|
||||
|
||||
class Qwen25_7BVLI(BaseLlama, nn.Module):
|
||||
def __init__(self, config_dict):
|
||||
super().__init__()
|
||||
|
||||
@ -1,4 +1,4 @@
|
||||
# https://github.com/comfyanonymous/ComfyUI/blob/v0.3.64/comfy/ldm/lumina/model.py
|
||||
# https://github.com/comfyanonymous/ComfyUI/blob/v0.3.75/comfy/ldm/lumina/model.py
|
||||
# Reference: https://github.com/Alpha-VLLM/Lumina-Image-2.0
|
||||
|
||||
import math
|
||||
@ -39,12 +39,12 @@ def timestep_embedding(timesteps, dim, max_period=10000, repeat_only=False):
|
||||
|
||||
class TimestepEmbedder(nn.Module):
|
||||
|
||||
def __init__(self, hidden_size, frequency_embedding_size=256):
|
||||
def __init__(self, hidden_size, frequency_embedding_size=256, output_size=None):
|
||||
super().__init__()
|
||||
self.mlp = nn.Sequential(
|
||||
nn.Linear(frequency_embedding_size, hidden_size, bias=True),
|
||||
nn.SiLU(),
|
||||
nn.Linear(hidden_size, hidden_size, bias=True),
|
||||
nn.Linear(hidden_size, output_size or hidden_size, bias=True),
|
||||
)
|
||||
self.frequency_embedding_size = frequency_embedding_size
|
||||
|
||||
@ -56,7 +56,7 @@ class TimestepEmbedder(nn.Module):
|
||||
|
||||
class JointAttention(nn.Module):
|
||||
|
||||
def __init__(self, dim: int, n_heads: int, n_kv_heads: int, qk_norm: bool):
|
||||
def __init__(self, dim: int, n_heads: int, n_kv_heads: int, qk_norm: bool, out_bias: bool = False):
|
||||
super().__init__()
|
||||
self.n_kv_heads = n_heads if n_kv_heads is None else n_kv_heads
|
||||
self.n_local_heads = n_heads
|
||||
@ -65,7 +65,7 @@ class JointAttention(nn.Module):
|
||||
self.head_dim = dim // n_heads
|
||||
|
||||
self.qkv = nn.Linear(dim, (n_heads + self.n_kv_heads + self.n_kv_heads) * self.head_dim, bias=False)
|
||||
self.out = nn.Linear(n_heads * self.head_dim, dim, bias=False)
|
||||
self.out = nn.Linear(n_heads * self.head_dim, dim, bias=out_bias)
|
||||
|
||||
if qk_norm:
|
||||
self.q_norm = nn.RMSNorm(self.head_dim, elementwise_affine=True)
|
||||
@ -129,14 +129,14 @@ class FeedForward(nn.Module):
|
||||
|
||||
|
||||
class JointTransformerBlock(nn.Module):
|
||||
def __init__(self, layer_id: int, dim: int, n_heads: int, n_kv_heads: int, multiple_of: int, ffn_dim_multiplier: float, norm_eps: float, qk_norm: bool, modulation=True):
|
||||
def __init__(self, layer_id: int, dim: int, n_heads: int, n_kv_heads: int, multiple_of: int, ffn_dim_multiplier: float, norm_eps: float, qk_norm: bool, modulation=True, z_image_modulation=False, attn_out_bias=False):
|
||||
super().__init__()
|
||||
self.dim = dim
|
||||
self.head_dim = dim // n_heads
|
||||
self.attention = JointAttention(dim, n_heads, n_kv_heads, qk_norm)
|
||||
self.attention = JointAttention(dim, n_heads, n_kv_heads, qk_norm, out_bias=attn_out_bias)
|
||||
self.feed_forward = FeedForward(
|
||||
dim=dim,
|
||||
hidden_dim=4 * dim,
|
||||
hidden_dim=dim,
|
||||
multiple_of=multiple_of,
|
||||
ffn_dim_multiplier=ffn_dim_multiplier,
|
||||
)
|
||||
@ -149,10 +149,15 @@ class JointTransformerBlock(nn.Module):
|
||||
|
||||
self.modulation = modulation
|
||||
if modulation:
|
||||
self.adaLN_modulation = nn.Sequential(
|
||||
nn.SiLU(),
|
||||
nn.Linear(min(dim, 1024), 4 * dim, bias=True),
|
||||
)
|
||||
if z_image_modulation:
|
||||
self.adaLN_modulation = nn.Sequential(
|
||||
nn.Linear(min(dim, 256), 4 * dim, bias=True),
|
||||
)
|
||||
else:
|
||||
self.adaLN_modulation = nn.Sequential(
|
||||
nn.SiLU(),
|
||||
nn.Linear(min(dim, 1024), 4 * dim, bias=True),
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor, x_mask: torch.Tensor, freqs_cis: torch.Tensor, adaln_input: torch.Tensor = None, transformer_options={}):
|
||||
if self.modulation:
|
||||
@ -192,14 +197,19 @@ class JointTransformerBlock(nn.Module):
|
||||
|
||||
class FinalLayer(nn.Module):
|
||||
|
||||
def __init__(self, hidden_size, patch_size, out_channels):
|
||||
def __init__(self, hidden_size, patch_size, out_channels, z_image_modulation=False):
|
||||
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)
|
||||
|
||||
if z_image_modulation:
|
||||
min_mod = 256
|
||||
else:
|
||||
min_mod = 1024
|
||||
|
||||
self.adaLN_modulation = nn.Sequential(
|
||||
nn.SiLU(),
|
||||
nn.Linear(min(hidden_size, 1024), hidden_size, bias=True),
|
||||
nn.Linear(min(hidden_size, min_mod), hidden_size, bias=True),
|
||||
)
|
||||
|
||||
def forward(self, x, c):
|
||||
@ -221,18 +231,24 @@ class NextDiT(nn.Module):
|
||||
n_heads: int = 32,
|
||||
n_kv_heads: int = None,
|
||||
multiple_of: int = 256,
|
||||
ffn_dim_multiplier: float = None,
|
||||
ffn_dim_multiplier: float = 4.0,
|
||||
norm_eps: float = 1e-5,
|
||||
qk_norm: bool = False,
|
||||
cap_feat_dim: int = 5120,
|
||||
axes_dims: list[int] = (16, 56, 56),
|
||||
axes_lens: list[int] = (1, 512, 512),
|
||||
rope_theta: float = 10000.0,
|
||||
z_image_modulation: bool = False,
|
||||
time_scale: float = 1.0,
|
||||
pad_tokens_multiple: int = None,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.in_channels = in_channels
|
||||
self.out_channels = in_channels
|
||||
self.patch_size = patch_size
|
||||
self.time_scale = time_scale
|
||||
self.pad_tokens_multiple = pad_tokens_multiple
|
||||
|
||||
self.x_embedder = nn.Linear(in_features=patch_size * patch_size * in_channels, out_features=dim, bias=True)
|
||||
|
||||
@ -248,6 +264,7 @@ class NextDiT(nn.Module):
|
||||
norm_eps,
|
||||
qk_norm,
|
||||
modulation=True,
|
||||
z_image_modulation=z_image_modulation,
|
||||
)
|
||||
for layer_id in range(n_refiner_layers)
|
||||
]
|
||||
@ -269,7 +286,7 @@ class NextDiT(nn.Module):
|
||||
]
|
||||
)
|
||||
|
||||
self.t_embedder = TimestepEmbedder(min(dim, 1024))
|
||||
self.t_embedder = TimestepEmbedder(min(dim, 1024), output_size=256 if z_image_modulation else None)
|
||||
self.cap_embedder = nn.Sequential(
|
||||
nn.RMSNorm(cap_feat_dim, eps=norm_eps, elementwise_affine=True),
|
||||
nn.Linear(cap_feat_dim, dim, bias=True),
|
||||
@ -286,17 +303,23 @@ class NextDiT(nn.Module):
|
||||
ffn_dim_multiplier,
|
||||
norm_eps,
|
||||
qk_norm,
|
||||
z_image_modulation=z_image_modulation,
|
||||
attn_out_bias=False,
|
||||
)
|
||||
for layer_id in range(n_layers)
|
||||
]
|
||||
)
|
||||
self.norm_final = nn.RMSNorm(dim, eps=norm_eps, elementwise_affine=True)
|
||||
self.final_layer = FinalLayer(dim, patch_size, self.out_channels)
|
||||
self.final_layer = FinalLayer(dim, patch_size, self.out_channels, z_image_modulation=z_image_modulation)
|
||||
|
||||
if self.pad_tokens_multiple is not None:
|
||||
self.x_pad_token = nn.Parameter(torch.empty((1, dim)))
|
||||
self.cap_pad_token = nn.Parameter(torch.empty((1, dim)))
|
||||
|
||||
assert (dim // n_heads) == sum(axes_dims)
|
||||
self.axes_dims = axes_dims
|
||||
self.axes_lens = axes_lens
|
||||
self.rope_embedder = EmbedND(dim=dim // n_heads, theta=10000.0, axes_dim=axes_dims)
|
||||
self.rope_embedder = EmbedND(dim=dim // n_heads, theta=rope_theta, axes_dim=axes_dims)
|
||||
self.dim = dim
|
||||
self.n_heads = n_heads
|
||||
|
||||
@ -319,96 +342,43 @@ class NextDiT(nn.Module):
|
||||
device = x[0].device
|
||||
dtype = x[0].dtype
|
||||
|
||||
if cap_mask is not None:
|
||||
l_effective_cap_len = cap_mask.sum(dim=1).tolist()
|
||||
else:
|
||||
l_effective_cap_len = [num_tokens] * bsz
|
||||
if self.pad_tokens_multiple is not None:
|
||||
pad_extra = (-cap_feats.shape[1]) % self.pad_tokens_multiple
|
||||
cap_feats = torch.cat((cap_feats, self.cap_pad_token.to(device=cap_feats.device, dtype=cap_feats.dtype, copy=True).unsqueeze(0).repeat(cap_feats.shape[0], pad_extra, 1)), dim=1)
|
||||
|
||||
if cap_mask is not None and not torch.is_floating_point(cap_mask):
|
||||
cap_mask = (cap_mask - 1).to(dtype) * torch.finfo(dtype).max
|
||||
cap_pos_ids = torch.zeros(bsz, cap_feats.shape[1], 3, dtype=torch.float32, device=device)
|
||||
cap_pos_ids[:, :, 0] = torch.arange(cap_feats.shape[1], dtype=torch.float32, device=device) + 1.0
|
||||
|
||||
img_sizes = [(img.size(1), img.size(2)) for img in x]
|
||||
l_effective_img_len = [(H // pH) * (W // pW) for (H, W) in img_sizes]
|
||||
B, C, H, W = x.shape
|
||||
x = self.x_embedder(x.view(B, C, H // pH, pH, W // pW, pW).permute(0, 2, 4, 3, 5, 1).flatten(3).flatten(1, 2))
|
||||
|
||||
max_seq_len = max((cap_len + img_len for cap_len, img_len in zip(l_effective_cap_len, l_effective_img_len)))
|
||||
max_cap_len = max(l_effective_cap_len)
|
||||
max_img_len = max(l_effective_img_len)
|
||||
H_tokens, W_tokens = H // pH, W // pW
|
||||
x_pos_ids = torch.zeros((bsz, x.shape[1], 3), dtype=torch.float32, device=device)
|
||||
x_pos_ids[:, :, 0] = cap_feats.shape[1] + 1
|
||||
x_pos_ids[:, :, 1] = torch.arange(H_tokens, dtype=torch.float32, device=device).view(-1, 1).repeat(1, W_tokens).flatten()
|
||||
x_pos_ids[:, :, 2] = torch.arange(W_tokens, dtype=torch.float32, device=device).view(1, -1).repeat(H_tokens, 1).flatten()
|
||||
|
||||
position_ids = torch.zeros(bsz, max_seq_len, 3, dtype=torch.int32, device=device)
|
||||
if self.pad_tokens_multiple is not None:
|
||||
pad_extra = (-x.shape[1]) % self.pad_tokens_multiple
|
||||
x = torch.cat((x, self.x_pad_token.to(device=x.device, dtype=dtype, copy=True).unsqueeze(0).repeat(x.shape[0], pad_extra, 1)), dim=1)
|
||||
x_pos_ids = torch.nn.functional.pad(x_pos_ids, (0, 0, 0, pad_extra))
|
||||
|
||||
for i in range(bsz):
|
||||
cap_len = l_effective_cap_len[i]
|
||||
img_len = l_effective_img_len[i]
|
||||
H, W = img_sizes[i]
|
||||
H_tokens, W_tokens = H // pH, W // pW
|
||||
assert H_tokens * W_tokens == img_len
|
||||
|
||||
position_ids[i, :cap_len, 0] = torch.arange(cap_len, dtype=torch.int32, device=device)
|
||||
position_ids[i, cap_len : cap_len + img_len, 0] = cap_len
|
||||
row_ids = torch.arange(H_tokens, dtype=torch.int32, device=device).view(-1, 1).repeat(1, W_tokens).flatten()
|
||||
col_ids = torch.arange(W_tokens, dtype=torch.int32, device=device).view(1, -1).repeat(H_tokens, 1).flatten()
|
||||
position_ids[i, cap_len : cap_len + img_len, 1] = row_ids
|
||||
position_ids[i, cap_len : cap_len + img_len, 2] = col_ids
|
||||
|
||||
freqs_cis = self.rope_embedder(position_ids).movedim(1, 2).to(dtype)
|
||||
|
||||
cap_freqs_cis_shape = list(freqs_cis.shape)
|
||||
cap_freqs_cis_shape[1] = cap_feats.shape[1]
|
||||
cap_freqs_cis = torch.zeros(*cap_freqs_cis_shape, device=device, dtype=freqs_cis.dtype)
|
||||
|
||||
img_freqs_cis_shape = list(freqs_cis.shape)
|
||||
img_freqs_cis_shape[1] = max_img_len
|
||||
img_freqs_cis = torch.zeros(*img_freqs_cis_shape, device=device, dtype=freqs_cis.dtype)
|
||||
|
||||
for i in range(bsz):
|
||||
cap_len = l_effective_cap_len[i]
|
||||
img_len = l_effective_img_len[i]
|
||||
cap_freqs_cis[i, :cap_len] = freqs_cis[i, :cap_len]
|
||||
img_freqs_cis[i, :img_len] = freqs_cis[i, cap_len : cap_len + img_len]
|
||||
freqs_cis = self.rope_embedder(torch.cat((cap_pos_ids, x_pos_ids), dim=1)).movedim(1, 2).to(dtype)
|
||||
|
||||
# refine context
|
||||
for layer in self.context_refiner:
|
||||
cap_feats = layer(cap_feats, cap_mask, cap_freqs_cis, transformer_options=transformer_options)
|
||||
cap_feats = layer(cap_feats, cap_mask, freqs_cis[:, : cap_pos_ids.shape[1]], transformer_options=transformer_options)
|
||||
|
||||
flat_x = []
|
||||
for i in range(bsz):
|
||||
img = x[i]
|
||||
C, H, W = img.size()
|
||||
img = img.view(C, H // pH, pH, W // pW, pW).permute(1, 3, 2, 4, 0).flatten(2).flatten(0, 1)
|
||||
flat_x.append(img)
|
||||
x = flat_x
|
||||
padded_img_embed = torch.zeros(bsz, max_img_len, x[0].shape[-1], device=device, dtype=x[0].dtype)
|
||||
padded_img_mask = torch.zeros(bsz, max_img_len, dtype=dtype, device=device)
|
||||
for i in range(bsz):
|
||||
padded_img_embed[i, : l_effective_img_len[i]] = x[i]
|
||||
padded_img_mask[i, l_effective_img_len[i] :] = -torch.finfo(dtype).max
|
||||
|
||||
padded_img_embed = self.x_embedder(padded_img_embed)
|
||||
padded_img_mask = padded_img_mask.unsqueeze(1)
|
||||
padded_img_mask = None
|
||||
for layer in self.noise_refiner:
|
||||
padded_img_embed = layer(padded_img_embed, padded_img_mask, img_freqs_cis, t, transformer_options=transformer_options)
|
||||
|
||||
if cap_mask is not None:
|
||||
mask = torch.zeros(bsz, max_seq_len, dtype=dtype, device=device)
|
||||
mask[:, :max_cap_len] = cap_mask[:, :max_cap_len]
|
||||
else:
|
||||
mask = None
|
||||
|
||||
padded_full_embed = torch.zeros(bsz, max_seq_len, self.dim, device=device, dtype=x[0].dtype)
|
||||
for i in range(bsz):
|
||||
cap_len = l_effective_cap_len[i]
|
||||
img_len = l_effective_img_len[i]
|
||||
|
||||
padded_full_embed[i, :cap_len] = cap_feats[i, :cap_len]
|
||||
padded_full_embed[i, cap_len : cap_len + img_len] = padded_img_embed[i, :img_len]
|
||||
x = layer(x, padded_img_mask, freqs_cis[:, cap_pos_ids.shape[1] :], t, transformer_options=transformer_options)
|
||||
|
||||
padded_full_embed = torch.cat((cap_feats, x), dim=1)
|
||||
mask = None
|
||||
img_sizes = [(H, W)] * bsz
|
||||
l_effective_cap_len = [cap_feats.shape[1]] * bsz
|
||||
return padded_full_embed, mask, img_sizes, l_effective_cap_len, freqs_cis
|
||||
|
||||
@staticmethod
|
||||
def load_tokens(transformer_options: dict) -> int:
|
||||
opt: list[int] = transformer_options.get("cond_or_uncond", [0])
|
||||
c_uc: str = "c" if opt[0] == 0 else "uc"
|
||||
return dynamic_args["num_tokens"][c_uc]
|
||||
|
||||
def forward(self, x, timesteps, context, num_tokens=None, attention_mask=None, **kwargs):
|
||||
t = 1.0 - timesteps
|
||||
cap_feats = context
|
||||
@ -416,13 +386,12 @@ class NextDiT(nn.Module):
|
||||
bs, c, h, w = x.shape
|
||||
x = pad_to_patch_size(x, (self.patch_size, self.patch_size))
|
||||
|
||||
t = self.t_embedder(t, dtype=x.dtype)
|
||||
t = self.t_embedder(t * self.time_scale, dtype=x.dtype) # (N, D)
|
||||
adaln_input = t
|
||||
|
||||
cap_feats = self.cap_embedder(cap_feats)
|
||||
cap_feats = self.cap_embedder(cap_feats) # (N, L, D) # todo check if able to batchify w.o. redundant compute
|
||||
|
||||
transformer_options = kwargs.get("transformer_options", {})
|
||||
num_tokens = num_tokens or self.load_tokens(transformer_options)
|
||||
|
||||
x_is_tensor = isinstance(x, torch.Tensor)
|
||||
x, mask, img_size, cap_size, freqs_cis = self.patchify_and_embed(x, cap_feats, cap_mask, t, num_tokens, transformer_options=transformer_options)
|
||||
|
||||
@ -9,7 +9,6 @@ if TYPE_CHECKING:
|
||||
import torch
|
||||
|
||||
from backend import memory_management
|
||||
from backend.args import dynamic_args
|
||||
from backend.text_processing import emphasis, parsing
|
||||
from modules.shared import opts
|
||||
|
||||
@ -114,7 +113,7 @@ class GemmaTextProcessingEngine:
|
||||
# tokens += [self.id_pad] * remaining_count
|
||||
# multipliers += [1.0] * remaining_count
|
||||
|
||||
z = self.process_tokens([tokens], [multipliers], texts.is_negative_prompt)[0]
|
||||
z = self.process_tokens([tokens], [multipliers])[0]
|
||||
line_z_values.append(z)
|
||||
cache[line] = line_z_values
|
||||
|
||||
@ -155,9 +154,8 @@ class GemmaTextProcessingEngine:
|
||||
|
||||
return torch.cat(embeds_out), torch.tensor(attention_masks, device=device, dtype=torch.long), num_tokens
|
||||
|
||||
def process_tokens(self, batch_tokens, batch_multipliers, negative: bool):
|
||||
def process_tokens(self, batch_tokens, batch_multipliers):
|
||||
embeds, mask, count = self.process_embeds(batch_tokens)
|
||||
dynamic_args["num_tokens"]["uc" if negative else "c"] = int(max(1, torch.sum(mask).item()))
|
||||
_, z = self.text_encoder(
|
||||
None,
|
||||
attention_mask=mask,
|
||||
|
||||
149
backend/text_processing/qwen3_engine.py
Normal file
149
backend/text_processing/qwen3_engine.py
Normal file
@ -0,0 +1,149 @@
|
||||
# https://github.com/comfyanonymous/ComfyUI/blob/v0.3.75/comfy/sd1_clip.py
|
||||
# https://github.com/comfyanonymous/ComfyUI/blob/v0.3.75/comfy/text_encoders/z_image.py
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from modules.prompt_parser import SdConditioning
|
||||
|
||||
import torch
|
||||
|
||||
from backend import memory_management
|
||||
from backend.text_processing import emphasis, parsing
|
||||
from modules.shared import opts
|
||||
|
||||
|
||||
class PromptChunk:
|
||||
def __init__(self):
|
||||
self.tokens = []
|
||||
self.multipliers = []
|
||||
|
||||
|
||||
class Qwen3TextProcessingEngine:
|
||||
def __init__(self, text_encoder, tokenizer):
|
||||
super().__init__()
|
||||
|
||||
self.text_encoder = text_encoder
|
||||
self.tokenizer = tokenizer
|
||||
|
||||
self.id_pad = 151643
|
||||
self.llama_template = "<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n"
|
||||
self.intermediate_output = -2
|
||||
self.layer_norm_hidden_state = False
|
||||
|
||||
def tokenize(self, texts):
|
||||
llama_texts = [self.llama_template.format(text) for text in texts]
|
||||
return self.tokenizer(llama_texts)["input_ids"]
|
||||
|
||||
def tokenize_line(self, line):
|
||||
parsed = parsing.parse_prompt_attention(line, self.emphasis.name)
|
||||
|
||||
tokenized = self.tokenize([text for text, _ in parsed])
|
||||
|
||||
chunks = []
|
||||
chunk = PromptChunk()
|
||||
token_count = 0
|
||||
|
||||
def next_chunk():
|
||||
nonlocal token_count
|
||||
nonlocal chunk
|
||||
|
||||
chunks.append(chunk)
|
||||
chunk = PromptChunk()
|
||||
|
||||
for tokens, (text, weight) in zip(tokenized, parsed):
|
||||
if text == "BREAK" and weight == -1:
|
||||
next_chunk()
|
||||
continue
|
||||
|
||||
position = 0
|
||||
while position < len(tokens):
|
||||
token = tokens[position]
|
||||
chunk.tokens.append(token)
|
||||
chunk.multipliers.append(weight)
|
||||
position += 1
|
||||
|
||||
if chunk.tokens or not chunks:
|
||||
next_chunk()
|
||||
|
||||
return chunks, token_count
|
||||
|
||||
def __call__(self, texts: "SdConditioning"):
|
||||
zs = []
|
||||
cache = {}
|
||||
|
||||
self.emphasis = emphasis.get_current_option(opts.emphasis)()
|
||||
|
||||
for line in texts:
|
||||
if line in cache:
|
||||
line_z_values = cache[line]
|
||||
else:
|
||||
chunks, token_count = self.tokenize_line(line)
|
||||
line_z_values = []
|
||||
|
||||
# pad all chunks to length of longest chunk
|
||||
# max_tokens = 0
|
||||
# for chunk in chunks:
|
||||
# max_tokens = max(len(chunk.tokens), max_tokens)
|
||||
|
||||
for chunk in chunks:
|
||||
tokens = chunk.tokens
|
||||
multipliers = chunk.multipliers
|
||||
|
||||
# remaining_count = max_tokens - len(tokens)
|
||||
# if remaining_count > 0:
|
||||
# tokens += [self.id_pad] * remaining_count
|
||||
# multipliers += [1.0] * remaining_count
|
||||
|
||||
z = self.process_tokens([tokens], [multipliers])[0]
|
||||
line_z_values.append(z)
|
||||
cache[line] = line_z_values
|
||||
|
||||
zs.extend(line_z_values)
|
||||
|
||||
return torch.stack(zs)
|
||||
|
||||
def process_embeds(self, batch_tokens):
|
||||
device = memory_management.text_encoder_device()
|
||||
self.text_encoder.to(device)
|
||||
|
||||
embeds_out = []
|
||||
attention_masks = []
|
||||
num_tokens = []
|
||||
|
||||
for tokens in batch_tokens:
|
||||
attention_mask = []
|
||||
tokens_temp = []
|
||||
eos = False
|
||||
index = 0
|
||||
|
||||
for t in tokens:
|
||||
token = int(t)
|
||||
attention_mask.append(0 if eos else 1)
|
||||
tokens_temp += [token]
|
||||
if not eos and token == self.id_pad:
|
||||
eos = True
|
||||
index += 1
|
||||
|
||||
tokens_embed = torch.tensor([tokens_temp], device=device, dtype=torch.long)
|
||||
tokens_embed = self.text_encoder.get_input_embeddings()(tokens_embed)
|
||||
|
||||
index = 0
|
||||
|
||||
embeds_out.append(tokens_embed)
|
||||
attention_masks.append(attention_mask)
|
||||
num_tokens.append(sum(attention_mask))
|
||||
|
||||
return torch.cat(embeds_out), torch.tensor(attention_masks, device=device, dtype=torch.long), num_tokens
|
||||
|
||||
def process_tokens(self, batch_tokens, batch_multipliers):
|
||||
embeds, mask, count = self.process_embeds(batch_tokens)
|
||||
_, z = self.text_encoder(
|
||||
None,
|
||||
attention_mask=mask,
|
||||
embeds=embeds,
|
||||
num_tokens=count,
|
||||
intermediate_output=self.intermediate_output,
|
||||
final_layer_norm_intermediate=self.layer_norm_hidden_state,
|
||||
)
|
||||
return z
|
||||
@ -7,7 +7,8 @@ from huggingface_hub import snapshot_download
|
||||
|
||||
PIPELINE_PATHS = (
|
||||
# "Qwen/Qwen-Image",
|
||||
"neta-art/Neta-Lumina-Diffusers",
|
||||
# "neta-art/Neta-Lumina-Diffusers",
|
||||
"Tongyi-MAI/Z-Image-Turbo",
|
||||
)
|
||||
|
||||
for pretrained in PIPELINE_PATHS:
|
||||
|
||||
@ -43,14 +43,31 @@ def detect_unet_config(state_dict: dict, key_prefix: str):
|
||||
dit_config["image_model"] = "lumina2"
|
||||
dit_config["patch_size"] = 2
|
||||
dit_config["in_channels"] = 16
|
||||
dit_config["dim"] = 2304
|
||||
dit_config["cap_feat_dim"] = int(state_dict["{}cap_embedder.1.weight".format(key_prefix)].shape[1])
|
||||
w = state_dict["{}cap_embedder.1.weight".format(key_prefix)]
|
||||
dit_config["dim"] = int(w.shape[0])
|
||||
dit_config["cap_feat_dim"] = int(w.shape[1])
|
||||
dit_config["n_layers"] = count_blocks(state_dict_keys, "{}layers.".format(key_prefix) + "{}.")
|
||||
dit_config["n_heads"] = 24
|
||||
dit_config["n_kv_heads"] = 8
|
||||
dit_config["qk_norm"] = True
|
||||
dit_config["axes_dims"] = [32, 32, 32]
|
||||
dit_config["axes_lens"] = [300, 512, 512]
|
||||
|
||||
if dit_config["dim"] == 2304: # Lumina 2
|
||||
dit_config["n_heads"] = 24
|
||||
dit_config["n_kv_heads"] = 8
|
||||
dit_config["axes_dims"] = [32, 32, 32]
|
||||
dit_config["axes_lens"] = [300, 512, 512]
|
||||
dit_config["rope_theta"] = 10000.0
|
||||
dit_config["ffn_dim_multiplier"] = 4.0
|
||||
elif dit_config["dim"] == 3840: # Z-Image
|
||||
dit_config["n_heads"] = 30
|
||||
dit_config["n_kv_heads"] = 30
|
||||
dit_config["axes_dims"] = [32, 48, 48]
|
||||
dit_config["axes_lens"] = [1536, 512, 512]
|
||||
dit_config["rope_theta"] = 256.0
|
||||
dit_config["ffn_dim_multiplier"] = 8.0 / 3.0
|
||||
dit_config["z_image_modulation"] = True
|
||||
dit_config["time_scale"] = 1000.0
|
||||
if "{}cap_pad_token".format(key_prefix) in state_dict_keys:
|
||||
dit_config["pad_tokens_multiple"] = 32
|
||||
|
||||
return dit_config
|
||||
|
||||
if "{}head.modulation".format(key_prefix) in state_dict_keys: # Wan 2.1
|
||||
|
||||
@ -319,6 +319,7 @@ class Lumina2(BASE):
|
||||
|
||||
unet_config = {
|
||||
"image_model": "lumina2",
|
||||
"dim": 2304,
|
||||
}
|
||||
|
||||
sampling_settings = {
|
||||
@ -351,6 +352,25 @@ class Lumina2(BASE):
|
||||
return {"gemma2_2b": "text_encoder"}
|
||||
|
||||
|
||||
class ZImage(Lumina2):
|
||||
huggingface_repo = "Tongyi-MAI/Z-Image-Turbo"
|
||||
|
||||
unet_config = {
|
||||
"image_model": "lumina2",
|
||||
"dim": 3840,
|
||||
}
|
||||
|
||||
sampling_settings = {
|
||||
"multiplier": 1.0,
|
||||
"shift": 3.0,
|
||||
}
|
||||
|
||||
memory_usage_factor = 1.7
|
||||
|
||||
def clip_target(self, state_dict={}):
|
||||
return {"qwen3_4b.transformer": "text_encoder"}
|
||||
|
||||
|
||||
class WAN21_T2V(BASE):
|
||||
huggingface_repo = "Wan-AI/Wan2.1-T2V-14B"
|
||||
|
||||
@ -444,6 +464,7 @@ models = [
|
||||
FluxSchnell,
|
||||
Chroma,
|
||||
Lumina2,
|
||||
ZImage,
|
||||
WAN21_T2V,
|
||||
WAN21_I2V,
|
||||
QwenImage,
|
||||
|
||||
Loading…
Reference in New Issue
Block a user