import os
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/combined_prompts_plain.txt"


num_inference_steps= 50

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/ldm3d'

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

# streetview and sd21
#pipe = StableDiffusionPipeline.from_single_file(streetview_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()

for i, prompt in enumerate(prompts, start=1):

    prompt = prompt.strip()
    sample_name = f"{i}_{prompt}_{HEIGHT}_{WIDTH}.png"

    save_path = os.path.join(output_directory, sample_name)

    # streetview and sd21
    #image = pipe(prompt=prompt, num_inference_steps=num_inference_steps, negative_prompt=n_prompt, height=HEIGHT, width=WIDTH).images[0]

    # ldm3d
    image = pipe(prompt=prompt, num_inference_steps=num_inference_steps, negative_prompt=n_prompt, height=HEIGHT, width=WIDTH).rgb[0]

    image.save(save_path)