Update app.py

app.py

@@ -86,9 +86,9 @@ class MagicTimeController:
         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").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_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")