github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Phase 1

Browse Source
This commit is contained in:
Jaret Burkett
2026-08-27 11:53:08 -06:00
parent e8d9cf6d35
commit 8db198ec0a
51 changed files with 985 additions and 1273 deletions
Download Patch File Download Diff File Expand all files Collapse all files

33
extensions_built_in/diffusion_models/zeta_chroma/zeta_chroma_model.py
View File

@@ -13,14 +13,13 @@ from toolkit.samplers.custom_flowmatch_sampler import (
CustomFlowMatchEulerDiscreteScheduler,
)
from toolkit.accelerator import unwrap_model
from optimum.quanto import freeze
from toolkit.util.quantize import quantize, get_qtype, quantize_model
from toolkit.util.quantize import quantize_model
from toolkit.memory_management import MemoryManager
from safetensors.torch import load_file
from optimum.quanto import QTensor
from toolkit.metadata import get_meta_for_safetensors
from safetensors.torch import load_file, save_file
from transformers import AutoTokenizer, Qwen3ForCausalLM
from toolkit.models.v2.text_encoders.qwen3 import Qwen3TextEncoder
from diffusers import AutoencoderKL
from toolkit.models.FakeVAE import FakeVAE
from .zeta_chroma_transformer import ZImageDCT, ZImageDCTParams, vae_flatten, vae_unflatten, prepare_latent_image_ids, make_text_position_ids
@@ -142,31 +141,9 @@ class ZetaChromaModel(BaseModel):
flush()
self.print_and_status_update("Text Encoder")
tokenizer = AutoTokenizer.from_pretrained(
base_model_path, subfolder="tokenizer", torch_dtype=dtype
)
text_encoder = Qwen3ForCausalLM.from_pretrained(
base_model_path, subfolder="text_encoder", torch_dtype=dtype
)
if (
self.model_config.layer_offloading
and self.model_config.layer_offloading_text_encoder_percent > 0
):
MemoryManager.attach(
text_encoder,
self.device_torch,
offload_percent=self.model_config.layer_offloading_text_encoder_percent,
)
text_encoder.to(self.device_torch, dtype=dtype)
flush()
if self.model_config.quantize_te:
self.print_and_status_update("Quantizing Text Encoder")
quantize(text_encoder, weights=get_qtype(self.model_config.qtype_te))
freeze(text_encoder)
flush()
tokenizer = Qwen3TextEncoder.load_tokenizer(base_model_path)
text_encoder = Qwen3TextEncoder.load_model(base_model_path, dtype=dtype)
self.prepare_text_encoder(text_encoder, dtype=dtype)
self.print_and_status_update("Loading VAE")
vae = FakeVAE(scaling_factor=1.0)