#!/usr/bin/env python3
from diffusers import DiffusionPipeline, DDPMScheduler
import torch
import time
import os
from pathlib import Path
from huggingface_hub import HfApi
import random
import numpy as np
from deepfloyd_if.modules import IFStageI, IFStageII, IFStageIII, T5Embedder
import sys

api = HfApi()
start_time = time.time()
seed = 0
use_diffusers = bool(int(sys.argv[1]))

t5_pos_embeds = torch.load("/home/patrick/tensors/embeds_orig.pt").to("cuda")
t5_neg_embeds = torch.load("/home/patrick/tensors/neg_embeds.pt").to("cuda")

def seed_everything(seed=None):
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = True
    return seed

if use_diffusers:
    pipe = DiffusionPipeline.from_pretrained("/home/patrick/if-diff-ckpts/IF-I-IF-v1.0", torch_dtype=torch.float32, use_safetensors=True, text_encoder=None, safety_checker=None)
    config = dict(pipe.scheduler.config)
    config["timestep_spacing"] = "even_border"
    pipe.scheduler = DDPMScheduler.from_config(config)
    pipe.to("cuda")

    with torch.no_grad():
        # text_embeddings = t5.get_text_embeddings([prompt])
        seed_everything(0)
        out_image = pipe(prompt_embeds=t5_pos_embeds, negative_prompt_embeds=t5_neg_embeds, num_inference_steps=5).images[0]
        out_image.save("/home/patrick/images/if_diff.png")
else:
    if_I = IFStageI(device="cuda", dir_or_name="/home/patrick/IF-I-IF-v1.0/", model_kwargs={"precision": "fp32"})
    if_I_kwargs = {}
    if_I_kwargs['negative_t5_embs'] = t5_neg_embeds
    if_I_kwargs['seed'] = seed
    if_I_kwargs['t5_embs'] = t5_pos_embeds
    if_I_kwargs['aspect_ratio'] = "1:1"
    if_I_kwargs['progress'] = True
    if_I_kwargs['sample_timestep_respacing'] = '5'

    seed_everything(0)
    stageI_generations, _ = if_I.embeddings_to_image(**if_I_kwargs)

    if_I.to_images(stageI_generations)[0].save("/home/patrick/images/if_ref.png")