skroed
/

audio_ldm2_handler

Model card Files Files and versions

skroed commited on Nov 9, 2023

Commit

ff2924b

·

1 Parent(s): b46ce34

Fix: to cuda later.

Files changed (1) hide show

handler.py +2 -1

handler.py CHANGED Viewed

@@ -9,7 +9,8 @@ class EndpointHandler:
         # load model and processor from path
         self.pipeline = AudioLDM2Pipeline.from_pretrained(
             "cvssp/audioldm2-music", torch_dtype=torch.float16
-        ).to("cuda")
         self.pipeline.unet = torch.compile(
             self.pipeline.unet, mode="reduce-overhead", fullgraph=True
         )

         # load model and processor from path
         self.pipeline = AudioLDM2Pipeline.from_pretrained(
             "cvssp/audioldm2-music", torch_dtype=torch.float16
+        )
+        self.pipeline.to("cuda")
         self.pipeline.unet = torch.compile(
             self.pipeline.unet, mode="reduce-overhead", fullgraph=True
         )