wehub-resource-sync
64 lines
2.3 KiB
Python
64 lines
2.3 KiB
Python
import os
|
|
import tempfile
|
|
|
|
import torch
|
|
from diffusers.models.autoencoders.autoencoder_kl_flux2 import AutoencoderKLFlux2
|
|
from diffusers.models.transformers.transformer_flux2 import Flux2Transformer2DModel
|
|
from diffusers.pipelines.flux2.pipeline_flux2_klein import Flux2KleinPipeline
|
|
from diffusers.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler
|
|
from transformers import AutoTokenizer, Qwen3Config, Qwen3ForCausalLM
|
|
|
|
TINY_MODEL_DIR = os.path.join(tempfile.gettempdir(), "vllm-omni-tiny-models")
|
|
|
|
|
|
def _get_tiny_model_path(name: str) -> str:
|
|
path = os.path.join(TINY_MODEL_DIR, name)
|
|
os.makedirs(path, exist_ok=True)
|
|
return path
|
|
|
|
|
|
def tiny_flux2_klein_builder() -> str:
|
|
"""Build a tiny Flux2Klein model."""
|
|
model_dir = _get_tiny_model_path("Flux2KleinPipeline")
|
|
|
|
pipe = Flux2KleinPipeline(
|
|
scheduler=FlowMatchEulerDiscreteScheduler(),
|
|
vae=AutoencoderKLFlux2(
|
|
down_block_types=("DownEncoderBlock2D",),
|
|
up_block_types=("UpDecoderBlock2D",),
|
|
block_out_channels=(32,),
|
|
layers_per_block=1,
|
|
latent_channels=16,
|
|
norm_num_groups=16,
|
|
),
|
|
# NOTE: For now we need 28 layers because of hardcoded stuff in the model :(
|
|
text_encoder=Qwen3ForCausalLM(
|
|
Qwen3Config(
|
|
hidden_size=32,
|
|
intermediate_size=64,
|
|
num_hidden_layers=28,
|
|
num_attention_heads=2,
|
|
num_key_value_heads=2,
|
|
head_dim=16,
|
|
vocab_size=151936,
|
|
max_position_embeddings=512,
|
|
)
|
|
),
|
|
tokenizer=AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B"),
|
|
# NOTE: For now we need at least 2 layers for the transformer
|
|
# due to hardcoded hacks in CacheDiT for Flux2Klein specifically.
|
|
transformer=Flux2Transformer2DModel(
|
|
in_channels=64,
|
|
num_layers=2,
|
|
num_single_layers=2,
|
|
attention_head_dim=16,
|
|
num_attention_heads=2,
|
|
joint_attention_dim=96,
|
|
timestep_guidance_channels=32,
|
|
axes_dims_rope=(4, 4, 4, 4),
|
|
),
|
|
)
|
|
# Need dtypes to be consistent; for now we just put it on bfloat16
|
|
pipe.to(torch.bfloat16).save_pretrained(model_dir)
|
|
return model_dir
|