"""
Reproduce: same seed, same model, same prompt — cpu vs cuda generator produces different images.
"""
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "1"

import sys
sys.path.insert(0, "Normalized-Attention-Guidance")
import torch
from diffusers import FluxPipeline

SEED = 134
PROMPT = "A vintage red bicycle leaning against a brick wall, its pedals leading back to a solitary rear gear."

print("Loading FluxPipeline...")
pipe = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16
).to("cuda")

kwargs = dict(
    guidance_scale=0.0,
    num_inference_steps=4,
    max_sequence_length=256,
)

img_cpu = pipe(PROMPT, generator=torch.Generator("cpu").manual_seed(SEED), **kwargs).images[0]
img_cpu.save("gen_cpu.png")
print("Saved gen_cpu.png")

img_cuda = pipe(PROMPT, generator=torch.Generator("cuda").manual_seed(SEED), **kwargs).images[0]
img_cuda.save("gen_cuda.png")
print("Saved gen_cuda.png")

# Check if they're identical
import numpy as np
arr_cpu = np.array(img_cpu)
arr_cuda = np.array(img_cuda)
print(f"Images identical: {np.array_equal(arr_cpu, arr_cuda)}")
print(f"Max pixel diff: {np.abs(arr_cpu.astype(int) - arr_cuda.astype(int)).max()}")