Save the following code to a Python file
krea.py
:
import argparse
import time
import torch
from diffusers import FluxPipeline
from dfloat11 import DFloat11Model
# Parse command line arguments
parser = argparse.ArgumentParser(description="Generate images using FLUX.1-Krea-dev model")
parser.add_argument(
"--prompt", type=str, help="Text prompt for image generation",
default="An astronaut, helmet off, sits at a tiny table set on the tip of a crescent moon, sipping tea while gazing at a swirling galaxy in the distance. Stars twinkle around, casting a gentle glow on the lunar surface.",
)
parser.add_argument("--width", type=int, default=1024, help="Image width")
parser.add_argument("--height", type=int, default=1024, help="Image height")
parser.add_argument("--guidance_scale", type=float, default=4.5, help="Guidance scale for generation")
parser.add_argument("--save_file_name", type=str, default="flux-krea-dev.png", help="Output file name")
parser.add_argument("--cpu_offload", action="store_true", help="Enable DFloat11 CPU offloading")
args = parser.parse_args()
# Load the pipeline
pipe = FluxPipeline.from_pretrained(
"black-forest-labs/FLUX.1-Krea-dev",
torch_dtype=torch.bfloat16,
)
# Load DFloat11 model
DFloat11Model.from_pretrained(
"DFloat11/FLUX.1-Krea-dev-DF11",
bfloat16_model=pipe.transformer,
device="cpu",
cpu_offload=args.cpu_offload,
)
pipe.enable_model_cpu_offload()
start_time = time.time()
# Generate image
image = pipe(
args.prompt,
height=args.height,
width=args.width,
guidance_scale=args.guidance_scale,
).images[0]
end_time = time.time()
# Save the image
image.save(args.save_file_name)
# Print time and memory usage
print(f"Time taken: {end_time - start_time:.2f} seconds")
peak_memory = torch.cuda.max_memory_allocated()
print(f"Peak memory: {peak_memory / 1000 ** 3:.2f} GB")