import torch

from diffusers import StableDiffusion3Pipeline
from diffusers.callbacks import SD3CFGCutoffCallback


torch_dtype = torch.bfloat16

pipe = StableDiffusion3Pipeline.from_pretrained("stabilityai/stable-diffusion-3.5-medium", torch_dtype=torch_dtype)
pipe.enable_model_cpu_offload()

cfg_callback = SD3CFGCutoffCallback(cutoff_step_ratio=0.3)

prompt = "A capybara in a field of flowers"

image = pipe(
    prompt,
    num_inference_steps=28,
    callback_on_step_end=cfg_callback,
    generator=torch.Generator(device="cuda").manual_seed(42),
).images[0]

image.save("capybara_cfg_cutoff.png")