from diffusers import DiffusionPipeline, FluxTransformer2DModel, NunchakuConfig
from nunchaku.models.linear import SVDQW4A4Linear
from safetensors import safe_open
from huggingface_hub import hf_hub_download
import torch 
def modules_without_qweight(safetensors_path: str):
 no_qweight = set()
 with safe_open(safetensors_path, framework="pt", device="cpu") as f:
 for key in f.keys():
 if key.endswith(".weight"):
 # module name is everything except the last piece after "."
 module_name = ".".join(key.split(".")[:-1])
 no_qweight.add(module_name)
 return sorted(no_qweight)
ckpt_id = "black-forest-labs/FLUX.1-dev"
state_dict_path = hf_hub_download(repo_id="nunchaku-tech/nunchaku-flux.1-dev", filename="svdq-int4_r32-flux.1-dev.safetensors")
modules_to_not_convert = modules_without_qweight(state_dict_path)
# print(f"{modules_to_convert=}")
model = FluxTransformer2DModel.from_single_file(
 state_dict_path,
 config=ckpt_id, 
 subfolder="transformer",
 torch_dtype=torch.bfloat16, 
 quantization_config=NunchakuConfig(
 weight_dtype="int4",
 weight_group_size=64,
 activation_dtype="int4",
 activation_group_size=64,
 modules_to_not_convert=modules_to_not_convert
 )
).to("cuda")
has_svd = any(isinstance(module, SVDQW4A4Linear) for _, module in model.named_modules())
assert has_svd
pipe = DiffusionPipeline.from_pretrained(
 ckpt_id, transformer=model, torch_dtype=torch.bfloat16
).to("cuda")
image = pipe(
 "A cat holding a sign that says hello world", 
 num_inference_steps=50, 
 guidance_scale=3.5,
 generator=torch.manual_seed(0),
).images[0]
image.save(f"nunchaku.png")