Views
No views yet
1"""
2Some bits are from https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_utils.py
3"""
4
5from huggingface_hub import hf_hub_download
6from accelerate.utils import set_module_tensor_to_device, compute_module_sizes
7from accelerate import init_empty_weights
8from convert_nf4_flux import _replace_with_bnb_linear, create_quantized_param, check_quantized_param
9from diffusers import FluxTransformer2DModel, FluxPipeline
10import safetensors.torch
11import gc
12import torch
13
14dtype = torch.bfloat16
15is_torch_e4m3fn_available = hasattr(torch, "float8_e4m3fn")
16ckpt_path = hf_hub_download("shauray/flux.1-dev-uncensored-nf4", filename="diffusion_pytorch_model.safetensors")
17original_state_dict = safetensors.torch.load_file(ckpt_path)
18
19with init_empty_weights():
20 config = FluxTransformer2DModel.load_config("shauray/flux.1-dev-uncensored-nf4")
21 model = FluxTransformer2DModel.from_config(config).to(dtype)
22 expected_state_dict_keys = list(model.state_dict().keys())
23
24_replace_with_bnb_linear(model, "nf4")
25
26for param_name, param in original_state_dict.items():
27 if param_name not in expected_state_dict_keys:
28 continue
29
30 is_param_float8_e4m3fn = is_torch_e4m3fn_available and param.dtype == torch.float8_e4m3fn
31 if torch.is_floating_point(param) and not is_param_float8_e4m3fn:
32 param = param.to(dtype)
33
34 if not check_quantized_param(model, param_name):
35 set_module_tensor_to_device(model, param_name, device=0, value=param)
36 else:
37 create_quantized_param(
38 model, param, param_name, target_device=0, state_dict=original_state_dict, pre_quantized=True
39 )
40
41del original_state_dict
42gc.collect()
43
44print(compute_module_sizes(model)[""] / 1024 / 1204)
45
46pipe = FluxPipeline.from_pretrained("black-forest-labs/flux.1-dev", transformer=model, torch_dtype=dtype)
47pipe.enable_model_cpu_offload()
48
49prompt = "A mystic cat with a sign that says hello world!"
50image = pipe(prompt, guidance_scale=3.5, num_inference_steps=50, generator=torch.manual_seed(0)).images[0]
51image.save("flux-nf4-dev-loaded.png")