Views
No views yet
diffusers library on my 16GB GPU, but it's difficult to train with flux-dev-fp8, so I want to use 4-bit weights to save VRAM.diffusers model load the quantized weights properly.
1git clone https://github.com/HighCWu/flux-4bit
2cd flux-4bitpip install -r requirements.txt1import torch
2
3from model import T5EncoderModel, FluxTransformer2DModel
4from diffusers import FluxPipeline
5
6
7text_encoder_2: T5EncoderModel = T5EncoderModel.from_pretrained(
8 "HighCWu/FLUX.1-dev-4bit",
9 subfolder="text_encoder_2",
10 torch_dtype=torch.bfloat16,
11 # hqq_4bit_compute_dtype=torch.float32,
12)
13
14transformer: FluxTransformer2DModel = FluxTransformer2DModel.from_pretrained(
15 "HighCWu/FLUX.1-dev-4bit",
16 subfolder="transformer",
17 torch_dtype=torch.bfloat16,
18)
19
20pipe: FluxPipeline = FluxPipeline.from_pretrained(
21 "black-forest-labs/FLUX.1-dev",
22 text_encoder_2=text_encoder_2,
23 transformer=transformer,
24 torch_dtype=torch.bfloat16,
25)
26pipe.enable_model_cpu_offload() # with cpu offload, it cost 8.5GB vram
27# pipe.remove_all_hooks()
28# pipe = pipe.to('cuda') # without cpu offload, it cost 11GB vram
29
30prompt = "realistic, best quality, extremely detailed, ray tracing, photorealistic, A blue cat holding a sign that says hello world"
31image = pipe(
32 prompt,
33 height=1024,
34 width=1024,
35 guidance_scale=3.5,
36 output_type="pil",
37 num_inference_steps=16,
38 max_sequence_length=512,
39 generator=torch.Generator("cpu").manual_seed(0)
40).images[0]
41image.show()
42