VideoStory/main.py
2025-04-20 11:41:31 +02:00

66 lines
1.5 KiB
Python

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("")