import torch
from diffusers import FluxPipeline
from diffusers.schedulers import FlowMatchStochasticOvershootDiscreteScheduler, FlowMatchStochasticCurvedEulerScheduler

pipe = FluxPipeline.from_pretrained('black-forest-labs/FLUX.1-Krea-dev', torch_dtype=torch.bfloat16)
pipe.to('cuda')

euler_scheduler = pipe.scheduler
overshoot_scheduler = FlowMatchStochasticOvershootDiscreteScheduler.from_pretrained('black-forest-labs/FLUX.1-Krea-dev', subfolder='scheduler')
cec_scheduler = FlowMatchStochasticCurvedEulerScheduler.from_pretrained('black-forest-labs/FLUX.1-Krea-dev', subfolder='scheduler')

prompt = "A kangaroo holding a sign that says 'Sting like a bee, float like a butterfly' wearing black sunglasses and boxing gloves, the golden gate bridge in the background."

pipe.scheduler = euler_scheduler
out = pipe(
    prompt=prompt,
    height=1024,
    width=1024,
    guidance_scale=3.5,
    num_inference_steps=28,
    max_sequence_length=512,
    generator=torch.Generator("cpu").manual_seed(0)
).images[0]
out.save('euler.jpg')

## CIRCULAR EULER
pipe.scheduler = cec_scheduler
out = pipe(
    prompt=prompt,
    height=1024,
    width=1024,
    guidance_scale=3.5,
    num_inference_steps=28,
    max_sequence_length=512,
    generator=torch.Generator("cpu").manual_seed(0)
).images[0]
out.save('circular_euler.jpg')

## OVERSHOOT C = 1.0 default
pipe.scheduler = overshoot_scheduler
pipe.scheduler.set_c(2.0)
out = pipe(
    prompt=prompt,
    height=1024,
    width=1024,
    guidance_scale=3.5,
    num_inference_steps=28,
    max_sequence_length=512,
    generator=torch.Generator("cpu").manual_seed(0)
).images[0]
out.save('overshoot.jpg')