Spaces:
Sleeping
Sleeping
| import torch | |
| import os | |
| os.environ['CURL_CA_BUNDLE'] = '' | |
| import argparse | |
| from omegaconf import OmegaConf | |
| from diffusers import DiffusionPipeline | |
| from vlogger.planning_utils.gpt4_utils import (ExtractProtagonist, | |
| ExtractAProtagonist, | |
| split_story, | |
| patch_story_scripts, | |
| refine_story_scripts, | |
| protagonist_place_reference1, | |
| translate_video_script, | |
| time_scripts, | |
| ) | |
| def main(args): | |
| story_path = args.story_path | |
| save_script_path = os.path.join(story_path.rsplit('/', 1)[0], "script") | |
| if not os.path.exists(save_script_path): | |
| os.makedirs(save_script_path) | |
| with open(story_path, "r") as story_file: | |
| story = story_file.read() | |
| # summerize protagonists and places | |
| protagonists_places_file_path = os.path.join(save_script_path, "protagonists_places.txt") | |
| if args.only_one_protagonist: | |
| character_places = ExtractAProtagonist(story, protagonists_places_file_path) | |
| else: | |
| character_places = ExtractProtagonist(story, protagonists_places_file_path) | |
| print("Protagonists and places OK", flush=True) | |
| # make script | |
| script_file_path = os.path.join(save_script_path, "video_prompts.txt") | |
| video_list = split_story(story, script_file_path) | |
| video_list = patch_story_scripts(story, video_list, script_file_path) | |
| video_list = refine_story_scripts(video_list, script_file_path) | |
| print("Scripts OK", flush=True) | |
| # think about the protagonist in each scene | |
| reference_file_path = os.path.join(save_script_path, "protagonist_place_reference.txt") | |
| reference_lists = protagonist_place_reference1(video_list, character_places, reference_file_path) | |
| print("Reference protagonist OK", flush=True) | |
| # translate the English script to Chinese | |
| zh_file_path = os.path.join(save_script_path, "zh_video_prompts.txt") | |
| zh_video_list = translate_video_script(video_list, zh_file_path) | |
| print("Translation OK", flush=True) | |
| # schedule the time of script | |
| time_file_path = os.path.join(save_script_path, "time_scripts.txt") | |
| time_list = time_scripts(video_list, time_file_path) | |
| print("Time script OK", flush=True) | |
| # make reference image | |
| base = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", | |
| torch_dtype=torch.float16, | |
| variant="fp16", | |
| use_safetensors=True, | |
| ).to("cuda") | |
| refiner = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-refiner-1.0", | |
| text_encoder_2=base.text_encoder_2, | |
| vae=base.vae, | |
| torch_dtype=torch.float16, | |
| use_safetensors=True, | |
| variant="fp16", | |
| ).to("cuda") | |
| ref_dir_path = os.path.join(story_path.rsplit('/', 1)[0], "ref_img") | |
| if not os.path.exists(ref_dir_path): | |
| os.makedirs(ref_dir_path) | |
| for key, value in character_places.items(): | |
| prompt = key + ", " + value | |
| img_path = os.path.join(ref_dir_path, key + ".jpg") | |
| image = base(prompt=prompt, | |
| output_type="latent", | |
| height=1024, | |
| width=1024, | |
| guidance_scale=7 | |
| ).images[0] | |
| image = refiner(prompt=prompt, image=image[None, :]).images[0] | |
| image.save(img_path) | |
| print("Reference image OK", flush=True) | |
| if __name__ == "__main__": | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument("--config", type=str, default="configs/vlog_write_script.yaml") | |
| args = parser.parse_args() | |
| omega_conf = OmegaConf.load(args.config) | |
| main(omega_conf) | |