Views
No views yet
diffusers, I am still investigating the issue with the ComfyUI implementation.diffusers library, ComfyUI, or any other model), although models that use architectures which are unfamiliar to me might be more difficult.diffusers1import torch
2from diffusers import Flux2KleinKVPipeline, Flux2Transformer2DModel
3from transformers.modeling_utils import no_init_weights
4from dfloat11 import DFloat11Model
5
6device = "cuda"
7dtype = torch.bfloat16
8model_path = "black-forest-labs/FLUX.2-klein-9b-kv"
9
10text_encoder = DFloat11Model.from_pretrained("DFloat11/Qwen3-8B-DF11", device="cpu")
11with no_init_weights():
12 transformer = Flux2Transformer2DModel.from_config(
13 Flux2Transformer2DModel.load_config(
14 "black-forest-labs/FLUX.2-klein-9b-kv", subfolder="transformer"
15 ),
16 torch_dtype=torch.bfloat16
17 ).to(torch.bfloat16)
18DFloat11Model.from_pretrained("mingyi456/FLUX.2-klein-9b-kv-DF11", device="cpu", bfloat16_model=transformer)
19pipe = Flux2KleinKVPipeline.from_pretrained("", text_encoder=text_encoder, transformer=transformer, torch_dtype=torch.bfloat16)
20pipe.to(device)
21
22# Text-to-image (no reference image)
23print("Generating text-to-image...")
24image = pipe(
25 prompt="A cat holding a sign that says hello world",
26 height=1024,
27 width=1024,
28 num_inference_steps=4,
29 generator=torch.Generator(device=device).manual_seed(0),
30).images[0]
31image.save("t2i_output.png")
32print("Saved t2i_output.png")
33
34# Image-to-image with KV cache (using the generated image as reference)
35print("Generating image-to-image with KV cache...")
36image_kv = pipe(
37 prompt="A cat dressed like a wizard",
38 image=image,
39 height=1024,
40 width=1024,
41 num_inference_steps=4,
42 generator=torch.Generator(device=device).manual_seed(0),
43).images[0]
44image_kv.save("kv_output.png")
45print("Saved kv_output.png")pattern_dict for compression:1pattern_dict = {
2 r"double_stream_modulation_img.linear": [],
3 r"double_stream_modulation_txt.linear": [],
4 r"single_stream_modulation.linear": [],
5 r"context_embedder": [],
6 r"transformer_blocks\.\d+" : (
7 "attn.to_q",
8 "attn.to_k",
9 "attn.to_v",
10 "attn.to_out.0",
11 "attn.add_q_proj",
12 "attn.add_k_proj",
13 "attn.add_v_proj",
14 "attn.to_add_out",
15 "ff.linear_in",
16 "ff.linear_out",
17 "ff_context.linear_in",
18 "ff_context.linear_out",
19 ),
20 r"single_transformer_blocks\.\d+" : (
21 "attn.to_qkv_mlp_proj",
22 "attn.to_out",
23 ),
24 r"norm_out.linear": [],
25}