Update app.py

app.py

@@ -15,7 +15,7 @@ from utils.unet import UNet3DConditionModel
 from utils.pipeline_magictime import MagicTimePipeline
 from utils.util import save_videos_grid, convert_ldm_unet_checkpoint, convert_ldm_clip_checkpoint, convert_ldm_vae_checkpoint, load_diffusers_lora_unet, convert_ldm_clip_text_model
 
-pretrained_model_path   = "runwayml/stable-diffusion-v1-5"
+pretrained_model_path   = "./ckpts/Base_Model/stable-diffusion-v1-5"
 inference_config_path   = "./sample_configs/RealisticVision.yaml"
 magic_adapter_s_path    = "./ckpts/Magic_Weights/magic_adapter_s/magic_adapter_s.ckpt"
 magic_adapter_t_path    = "./ckpts/Magic_Weights/magic_adapter_t"
@@ -63,7 +63,7 @@ os.system(f"rm -rf gradio_cached_examples/")
 
 
 class MagicTimeController:
-    def __init__(self):
+    def __init__(self, tokenizer, text_encoder, vae, unet, text_model):
         
         # config dirs
         self.basedir                = os.getcwd()
@@ -85,13 +85,18 @@ class MagicTimeController:
         # config models
         self.inference_config      = OmegaConf.load(inference_config_path)[1]
 
-        self.tokenizer             = CLIPTokenizer.from_pretrained(pretrained_model_path, subfolder="tokenizer")
-        self.text_encoder          = CLIPTextModel.from_pretrained(pretrained_model_path, subfolder="text_encoder").to('cuda')
-        self.vae                   = AutoencoderKL.from_pretrained(pretrained_model_path, subfolder="vae").to('cuda')
-        self.unet                  = UNet3DConditionModel.from_pretrained_2d(pretrained_model_path, subfolder="unet", unet_additional_kwargs=OmegaConf.to_container(self.inference_config.unet_additional_kwargs)).to('cuda')
-        
-        self.text_model = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14")
-        
+        # self.tokenizer             = CLIPTokenizer.from_pretrained(pretrained_model_path, subfolder="tokenizer")
+        # self.text_encoder          = CLIPTextModel.from_pretrained(pretrained_model_path, subfolder="text_encoder").cuda()
+        # self.vae                   = AutoencoderKL.from_pretrained(pretrained_model_path, subfolder="vae").cuda()
+        # self.unet                  = UNet3DConditionModel.from_pretrained_2d(pretrained_model_path, subfolder="unet", unet_additional_kwargs=OmegaConf.to_container(self.inference_config.unet_additional_kwargs)).cuda()
+        # self.text_model            = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14")
+
+        self.tokenizer    = tokenizer
+        self.text_encoder = text_encoder
+        self.vae          = vae
+        self.unet         = unet
+        self.text_model   = text_model
+
         self.update_motion_module(self.motion_module_list[0])
         self.update_dreambooth(self.dreambooth_list[0])
         
@@ -191,9 +196,14 @@ class MagicTimeController:
             "dreambooth": dreambooth_dropdown,
         }
         return gr.Video(value=save_sample_path), gr.Json(value=json_config)
-        
-controller = MagicTimeController()
 
+inference_config = OmegaConf.load(inference_config_path)[1]
+tokenizer        =  CLIPTokenizer.from_pretrained(pretrained_model_path, subfolder="tokenizer")
+text_encoder     =  CLIPTextModel.from_pretrained(pretrained_model_path, subfolder="text_encoder").cuda()
+vae              =  AutoencoderKL.from_pretrained(pretrained_model_path, subfolder="vae").cuda()
+unet             =  UNet3DConditionModel.from_pretrained_2d(pretrained_model_path, subfolder="unet", unet_additional_kwargs=OmegaConf.to_container(inference_config.unet_additional_kwargs)).cuda()
+text_model       =  CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14")
+controller       = MagicTimeController(tokenizer=tokenizer, text_encoder=text_encoder, vae=vae, unet=unet, text_model=text_model)
 
 def ui():
     with gr.Blocks(css=css) as demo: