from diffusers import Flux2KleinPipeline
import torch

device = "cuda"
dtype = torch.bfloat16

# 4b-base
repo_id = "black-forest-labs/FLUX.2-klein-base-4B"

pipeline = Flux2KleinPipeline.from_pretrained(repo_id, torch_dtype=dtype)
pipeline.enable_model_cpu_offload()

prompt = "a photo of a forest with mist swirling around the tree trunks. The word 'FLUX.2' is painted over it in big, red brush strokes with visible texture"
height = 768
width = 1360
num_inference_steps = 50
guidance_scale = 4.0
generator = torch.Generator(device=device).manual_seed(42)

out = pipeline(
    prompt=prompt,
    height=height,
    width=width,
    num_inference_steps=num_inference_steps,
    guidance_scale=guidance_scale,
    generator=generator
).images[0]

out.save("output_4b_base.png")

# 9b-base
repo_id = "black-forest-labs/FLUX.2-klein-base-9B"

pipeline = Flux2KleinPipeline.from_pretrained(repo_id, torch_dtype=dtype)
pipeline.enable_model_cpu_offload()

prompt = "a photo of a forest with mist swirling around the tree trunks. The word 'FLUX.2' is painted over it in big, red brush strokes with visible texture"
height = 768
width = 1360
num_inference_steps = 50
guidance_scale = 4.0
generator = torch.Generator(device=device).manual_seed(42)

out = pipeline(
    prompt=prompt,
    height=height,
    width=width,
    num_inference_steps=num_inference_steps,
    guidance_scale=guidance_scale,
    generator=generator
).images[0]

out.save("output_9b_base.png")


# 4b

repo_id = "black-forest-labs/FLUX.2-klein-4B"
pipeline = Flux2KleinPipeline.from_pretrained(repo_id, torch_dtype=dtype)
pipeline.enable_model_cpu_offload()
num_inference_steps = 4
guidance_scale = 1.0
generator = torch.Generator(device=device).manual_seed(42)

out = pipeline(
    prompt=prompt,
    height=height,
    width=width,
    num_inference_steps=num_inference_steps,
    guidance_scale=guidance_scale,
    generator=generator
).images[0]

out.save("output_4b.png")


# 9b
repo_id = "black-forest-labs/FLUX.2-klein-9B"
pipeline = Flux2KleinPipeline.from_pretrained(repo_id, torch_dtype=dtype)
pipeline.enable_model_cpu_offload()
num_inference_steps = 4
guidance_scale = 1.0
generator = torch.Generator(device=device).manual_seed(42)

out = pipeline(
    prompt=prompt,
    height=height,
    width=width,
    num_inference_steps=num_inference_steps,
    guidance_scale=guidance_scale,
    generator=generator
).images[0]

out.save("output_9b.png")