from diffusers import AnimateDiffPipeline, MotionAdapter, EulerDiscreteScheduler
from safetensors.torch import load_file
from huggingface_hub import hf_hub_download
from transformers import set_seed
import torch
import time

SEED = 42
device = "cpu"
model_dtype = torch.float16
WARM_UP = 4
RUN = 4

set_seed(SEED)

print("\nLoading AnimateDiff-Lightning model...")
model_id = "ByteDance/AnimateDiff-Lightning"
step = 4
ckpt = f"animatediff_lightning_{step}step_diffusers.safetensors"
base = "emilianJR/epiCRealism"

adapter = MotionAdapter().to(device, model_dtype)
adapter.load_state_dict(load_file(hf_hub_download(model_id, ckpt), device=device))
pipe = AnimateDiffPipeline.from_pretrained(base, motion_adapter=adapter, torch_dtype=model_dtype).to(device)
pipe.scheduler = EulerDiscreteScheduler.from_config(
    pipe.scheduler.config, timestep_spacing="trailing", beta_schedule="linear"
)

def run_inference():
    set_seed(SEED)
    with torch.no_grad():
        output = pipe(
            prompt="An astronaut riding a green horse",
            guidance_scale=1.0,
            num_inference_steps=4,
        ).frames[0]
    return output

# Warm up
print(f"\nWarming up ({WARM_UP} iterations)...")
for i in range(WARM_UP):
    run_inference()
    print(f"  Warm-up {i+1}/{WARM_UP} done")

# Benchmark
print(f"\nRunning benchmark ({RUN} iterations)...")
elapsed_times = []
for i in range(RUN):
    start = time.perf_counter()
    output = run_inference()
    end = time.perf_counter()
    elapsed = (end - start) * 1000  # ms
    elapsed_times.append(elapsed)
    print(f"  Run {i+1}/{RUN}: {elapsed:.2f} ms")

# Statistics
avg_time = sum(elapsed_times) / len(elapsed_times)
min_time = min(elapsed_times)
max_time = max(elapsed_times)
print(f"\n{'='*50}")
print(f"Results ({RUN} runs):")
print(f"  Average: {avg_time:.2f} ms")
print(f"  Min:     {min_time:.2f} ms")
print(f"  Max:     {max_time:.2f} ms")
print(f"{'='*50}")