import gen
import torch
import video
import os


SYSTEMPROMPTT = ""
SYSTEMPROMPTI = ""
SDBAD = ""

with open("promptUtoT.txt", "r") as file:
    SYSTEMPROMPTT = file.readlines()[0]

with open("promptTtoI.txt", "r") as file:
    SYSTEMPROMPTI = file.readlines()[0]

def main(prompt):

    llm = gen.loadllama()

    raw_text = gen.generate_story(prompt, SYSTEMPROMPTT, model=llm)
    image_prompts = []
    raw_text = raw_text.split("\n")
    raw_text = [item for item in raw_text if item != ""]

    for i in range(len(raw_text)):
        promptimg = "Context:\n"
        for l in range(0, i - 1):
            promptimg += raw_text[l] + " "
        promptimg += "Current Scene:\n" + raw_text[i]
        image_prompts.append(gen.generate_story(promptimg, SYSTEMPROMPTI, model=llm).strip("\n"))

    llm, speaker = gen.loadtts()

    for i in range(len(raw_text)):
        try:
            gen.text_to_speech(raw_text[i], llm, speaker, i)
        except:
            pass

    del llm
    del speaker
    torch.cuda.empty_cache()

    llm = gen.loadsdxl()

    for elem in range(len(image_prompts)):
        gen.stableDiffusion(image_prompts[elem], SDBAD, llm, elem)

    del llm
    torch.cuda.empty_cache()
    """
    llm = gen.loadsvd()

    for elem in range(len([name for name in os.listdir('./images')])):
        gen.stableVideoDiffusion(llm, f"./images/{elem}.png", elem)

    del llm
    torch.cuda.empty_cache()
    """

    video.create_video(".", ".", "out.mp4")

    return None

main("")