update

app.py

@@ -22,6 +22,13 @@ magic_adapter_s_path    = "./ckpts/Magic_Weights/magic_adapter_s/magic_adapter_s
 magic_adapter_t_path    = "./ckpts/Magic_Weights/magic_adapter_t"
 magic_text_encoder_path = "./ckpts/Magic_Weights/magic_text_encoder"
 
+if torch.cuda.is_available():
+    device = "cuda"
+elif torch.backends.mps.is_available():
+    device = "mps"
+else:
+    device = "cpu"
+
 css = """
 .toolbutton {
     margin-buttom: 0em 0em 0em 0em;
@@ -87,9 +94,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(device)
+        self.vae                   = AutoencoderKL.from_pretrained(pretrained_model_path, subfolder="vae").to(device)
+        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)
         self.text_model            = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14")
 
         # self.tokenizer    = tokenizer
@@ -162,7 +169,7 @@ class MagicTimeController:
         pipeline = MagicTimePipeline(
             vae=self.vae, text_encoder=self.text_encoder, tokenizer=self.tokenizer, unet=self.unet,
             scheduler=DDIMScheduler(**OmegaConf.to_container(self.inference_config.noise_scheduler_kwargs))
-        ).to("cuda")
+        ).to(device)
         
         if int(seed_textbox) > 0: seed = int(seed_textbox)
         else: seed = random.randint(1, 1e16)
@@ -171,7 +178,7 @@ class MagicTimeController:
         assert seed == torch.initial_seed()
         print(f"### seed: {seed}")
         
-        generator = torch.Generator(device="cuda")
+        generator = torch.Generator(device=device)
         generator.manual_seed(seed)
         
         sample = pipeline(
@@ -256,4 +263,4 @@ def ui():
 if __name__ == "__main__":
     demo = ui()
     demo.queue(max_size=20)
-    demo.launch()
+    demo.launch()

requirements.txt

@@ -1,17 +1,17 @@
-torch==2.2.2
-torchvision==0.17.2
-torchaudio==2.2.2
-xformers==0.0.25.post1
+#torch==2.2.2
+#torchvision==0.17.2
+#torchaudio==2.2.2
+#xformers==0.0.25.post1
 imageio==2.27.0
 gdown
 einops
 omegaconf
 safetensors
 gradio
-triton
+#triton
 imageio[ffmpeg]
 imageio[pyav]
 ms-swift
 accelerate==0.28.0
 diffusers==0.11.1
-transformers==4.38.2
+transformers==4.38.2