Views
No views yet
1from huggingface_hub import hf_hub_download
2from accelerate.utils import set_module_tensor_to_device, compute_module_sizes
3from accelerate import init_empty_weights
4from diffusers.loaders.single_file_utils import convert_flux_transformer_checkpoint_to_diffusers
5from convert_nf4_flux import _replace_with_bnb_linear, create_quantized_param, check_quantized_param
6from diffusers import FluxTransformer2DModel, FluxPipeline
7import safetensors.torch
8import gc
9import torch
10
11dtype = torch.bfloat16
12ckpt_path = hf_hub_download("black-forest-labs/flux.1-schnell", filename="flux1-schnell.safetensors")
13original_state_dict = safetensors.torch.load_file(ckpt_path)
14converted_state_dict = convert_flux_transformer_checkpoint_to_diffusers(original_state_dict)
15
16del original_state_dict
17gc.collect()
18
19with init_empty_weights():
20 config = FluxTransformer2DModel.load_config("black-forest-labs/flux.1-schnell", subfolder="transformer")
21 model = FluxTransformer2DModel.from_config(config).to(dtype)
22
23_replace_with_bnb_linear(model, "nf4")
24for param_name, param in converted_state_dict.items():
25 param = param.to(dtype)
26 if not check_quantized_param(model, param_name):
27 set_module_tensor_to_device(model, param_name, device=0, value=param)
28 else:
29 create_quantized_param(model, param, param_name, target_device=0)
30
31del converted_state_dict
32gc.collect()
33
34print(compute_module_sizes(model)[""] / 1024 / 1204)
35
36pipe = FluxPipeline.from_pretrained("black-forest-labs/flux.1-schnell", transformer=model, torch_dtype=dtype)
37pipe.enable_model_cpu_offload()
38
39prompt = "A mystic cat with a sign that says hello world!"
40image = pipe(prompt, guidance_scale=3.5, num_inference_steps=4, generator=torch.manual_seed(0)).images[0]
41image.save("flux-nf4-schnell.png")
42
43model.push_to_hub("skimai/flux.1-schnell-nf4")