# test sana sprint
"""
python scripts/convert_sana_to_diffusers.py \
  --orig_ckpt_path /raid/.cache/huggingface/yiyi/models--JunsongChen--Sana_Sprint_1600M_1024px/snapshots/8ecfdc7e6269e5b065f5b2cf3fac9a2a1a778c6a/Sana_Sprint_1600M_1024px_36K.pth \
  --model_type SanaSprint_1600M_P1_D20 \
  --image_size 1024 \
  --dump_path /raid/yiyi/Sana-Sprint-yiyi \
  --save_full_pipeline \
  --scheduler_type scm
"""

from diffusers import SanaSprintPipeline
import torch

device = "cuda:0"
dtype = torch.bfloat16

repo = "/raid/yiyi/Sana-Sprint-yiyi"

pipeline = SanaSprintPipeline.from_pretrained(repo, torch_dtype=dtype)
pipeline.to(device)


prompt = "a tiny astronaut hatching from an egg on the moon"

image = pipeline(prompt=prompt, num_inference_steps=2).images[0]
image.save("test_out.png")

# test sana sprint
# (pipeline)
test_max_timesteps = [1.57080, 1.56830, 1.56580, 1.56454, 1.56246, 1.55830, 1.55413, 1.55080, 1.54580]
test_intermediate_timesteps = [None, 1.0, 1.1, 1.2, 1.3, 1.4]
test_num_inference_steps = [1,2,4]

# test_max_timesteps = [1.57080]
# test_intermediate_timesteps = [None]
# test_num_inference_steps = [1]

from diffusers import SanaSprintPipeline
import torch

device = "cuda:0"
dtype = torch.bfloat16
repo = "/raid/yiyi/Sana-Sprint-yiyi"

def run_sana(pipeline, num_inference_steps, max_timesteps, intermediate_timesteps):
    prompt = "a tiny astronaut hatching from an egg on the moon"
    generator = torch.Generator(device=device).manual_seed(123)
    test_name = f"num_inference_steps_{num_inference_steps}_max_timesteps_{max_timesteps}_intermediate_timesteps_{intermediate_timesteps}"
    print(f"--------------------------------")
    print(f"Running test:")
    print(f"num_inference_steps: {num_inference_steps}")
    print(f"max_timesteps: {max_timesteps}")
    print(f"intermediate_timesteps: {intermediate_timesteps}")
    try:
        image = pipeline(prompt=prompt, num_inference_steps=num_inference_steps, max_timesteps=max_timesteps, intermediate_timesteps=intermediate_timesteps, generator=generator).images[0]
        image.save(f"yiyi_test_10_1_out_{test_name}.png")
    except Exception as e:
        print(e)
    print(f"--------------------------------")


pipeline = SanaSprintPipeline.from_pretrained(repo, torch_dtype=dtype)
pipeline.to(device)

for num_inference_steps in test_num_inference_steps:
    for max_timesteps in test_max_timesteps:
        for intermediate_timesteps in test_intermediate_timesteps:
            run_sana(pipeline, num_inference_steps, max_timesteps, intermediate_timesteps)