import sys
import torch
from diffusers.schedulers.scheduling_dpmsolver_multistep import \
    DPMSolverMultistepScheduler
from diffusers.schedulers.scheduling_unipc_multistep import \
    UniPCMultistepScheduler

from diffusers import FluxPipeline


def main(scheduler):
    ckpt = "black-forest-labs/FLUX.1-dev"
    pipe = FluxPipeline.from_pretrained(ckpt, torch_dtype=torch.bfloat16).to("cuda")

    print("original scheduler", pipe.scheduler)

    if scheduler == 'dpm':
        pipe.scheduler = DPMSolverMultistepScheduler.from_config(
            pipe.scheduler.config, use_flow_sigmas=True, prediction_type="flow_prediction")
    elif scheduler == 'unipc':
        pipe.scheduler = UniPCMultistepScheduler.from_config(
            pipe.scheduler.config, use_flow_sigmas=True, prediction_type="flow_prediction")
    print("updated scheduler", pipe.scheduler)

    # prompt = "A cat holding a sign that says hello world"
    # prompt = "a white bird in front of a dinosaur standing by some trees"
    prompt = 'a motorcycle parked in an ornate bank lobby'
    num_inference_steps = 20
    generator = torch.Generator("cpu").manual_seed(0)
    image = pipe(
        prompt,
        height=1024,
        width=1024,
        guidance_scale=3.5,
        num_inference_steps=num_inference_steps,
        max_sequence_length=512,
        generator=generator,
    ).images[0]

    image.save(f"flux-diffusers-{num_inference_steps}.png")


if __name__ == '__main__':
    scheduler = sys.argv[1]
    main(scheduler)