|
|
import torch |
|
|
import os |
|
|
from dotenv import load_dotenv,find_dotenv |
|
|
load_dotenv(find_dotenv()) |
|
|
token=os.getenv('HF_TOKEN') |
|
|
from diffusers import FluxPipeline, FluxTransformer2DModel, GGUFQuantizationConfig |
|
|
|
|
|
ckpt_path = ( |
|
|
"https://huggingface.co/city96/FLUX.1-schnell-gguf/blob/main/flux1-schnell-Q2_K.gguf" |
|
|
|
|
|
) |
|
|
transformer = FluxTransformer2DModel.from_single_file( |
|
|
ckpt_path, |
|
|
quantization_config=GGUFQuantizationConfig(compute_dtype=torch.bfloat16), |
|
|
torch_dtype=torch.bfloat16, |
|
|
) |
|
|
pipe = FluxPipeline.from_pretrained( |
|
|
"black-forest-labs/FLUX.1-schnell", |
|
|
transformer=transformer, |
|
|
torch_dtype=torch.bfloat16, |
|
|
token=token, |
|
|
) |
|
|
pipe.enable_model_cpu_offload() |
|
|
prompt = "A cat holding a sign that says hello world" |
|
|
image = pipe(prompt, generator=torch.manual_seed(0)).images[0] |
|
|
image.save("flux-gguf.png") |
