import os
import torch
from datetime import datetime

from diffusers import StableDiffusionPipeline, StableDiffusionLDM3DPipeline

streetview_path="streetview-v4-10000-images.safetensors"

sd_diffusers_path = "checkpoints/v2-1_768-nonema-pruned.safetensors"

ldm3d_path = "checkpoints/ldm3d-pano"

prompt_file_path = "prompts/clip_score_prompts.txt"

num_inference_steps= 40

sample_basename = "streetview"

n_prompt = "bad, deformed, ugly, low quality, underexposed, overexposed, grainy, cartoon, anime"

original_config_file="v2-inference-v.yaml"

HEIGHT = 768
WIDTH = 2 * HEIGHT

timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")

output_directory = f'Generated Images/sd_clip'

# Create the directory if it does not exist
os.makedirs(output_directory, exist_ok=True)

# streetview and sd21

pipe = StableDiffusionPipeline.from_single_file(sd_diffusers_path, use_auth_token=False, local_files_only=True, original_config_file=original_config_file)

#ldm3d-pano

#pipe = StableDiffusionLDM3DPipeline.from_pretrained(ldm3d_path, use_auth_token=False, local_files_only=True, original_config_file=original_config_file)

pipe = pipe.to("cuda")

with open(prompt_file_path, 'r') as file:

    prompts = file.readlines()


IMAGES_PER_PROMPT = 5

generators = [torch.Generator(device='cuda').manual_seed(i) for i in range(IMAGES_PER_PROMPT)]


for i, prompt in enumerate(prompts):

    prompt = prompt.strip()

    for j in range(IMAGES_PER_PROMPT):

        sample_name = f"{i*IMAGES_PER_PROMPT + j}_{prompt}_{HEIGHT}_{WIDTH}.png"
        save_path = os.path.join(output_directory, sample_name)
        # streetview and sd21
        image = pipe(prompt=prompt, generator=generators[j], num_inference_steps=num_inference_steps, negative_prompt=n_prompt, height=HEIGHT, width=WIDTH).images[0]
        # ldm3d
        #image = pipe(prompt=prompt, generator=generators[j], num_inference_steps=num_inference_steps, negative_prompt=n_prompt, height=HEIGHT, width=WIDTH).rgb[0]
        image.save(save_path)