(self)
| 80 | |
| 81 | class MagicTimeController: |
| 82 | def __init__(self): |
| 83 | # config dirs |
| 84 | self.basedir = os.getcwd() |
| 85 | self.stable_diffusion_dir = os.path.join(self.basedir, model_path, "Base_Model") |
| 86 | self.motion_module_dir = os.path.join(self.basedir, model_path, "Base_Model", "motion_module") |
| 87 | self.personalized_model_dir = os.path.join(self.basedir, model_path, "DreamBooth") |
| 88 | self.savedir = os.path.join(self.basedir, "outputs") |
| 89 | os.makedirs(self.savedir, exist_ok=True) |
| 90 | |
| 91 | self.dreambooth_list = [] |
| 92 | self.motion_module_list = [] |
| 93 | |
| 94 | self.selected_dreambooth = None |
| 95 | self.selected_motion_module = None |
| 96 | |
| 97 | self.refresh_motion_module() |
| 98 | self.refresh_personalized_model() |
| 99 | |
| 100 | # config models |
| 101 | self.inference_config = OmegaConf.load(inference_config_path)[1] |
| 102 | |
| 103 | self.tokenizer = CLIPTokenizer.from_pretrained(pretrained_model_path, subfolder="tokenizer") |
| 104 | self.text_encoder = CLIPTextModel.from_pretrained(pretrained_model_path, subfolder="text_encoder").to(device) |
| 105 | self.vae = AutoencoderKL.from_pretrained(pretrained_model_path, subfolder="vae").to(device) |
| 106 | self.unet = UNet3DConditionModel.from_pretrained_2d(pretrained_model_path, subfolder="unet", unet_additional_kwargs=OmegaConf.to_container(self.inference_config.unet_additional_kwargs)).to(device) |
| 107 | self.text_model = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14") |
| 108 | self.unet_model = UNet3DConditionModel.from_pretrained_2d(pretrained_model_path, subfolder="unet", unet_additional_kwargs=OmegaConf.to_container(self.inference_config.unet_additional_kwargs)) |
| 109 | |
| 110 | self.update_motion_module(self.motion_module_list[0]) |
| 111 | self.update_motion_module_2(self.motion_module_list[0]) |
| 112 | self.update_dreambooth(self.dreambooth_list[0]) |
| 113 | |
| 114 | def refresh_motion_module(self): |
| 115 | motion_module_list = glob(os.path.join(self.motion_module_dir, "*.ckpt")) |
nothing calls this directly
no test coverage detected