envs

app.py

@@ -326,6 +326,7 @@ class ImageConductor:
         controlnet_images = rearrange(controlnet_images, "b f c h w -> b c f h w")
         num_controlnet_images = controlnet_images.shape[2]
         controlnet_images = rearrange(controlnet_images, "b c f h w -> (b f) c h w")
+        self.vae.to(device)
         controlnet_images = self.vae.encode(controlnet_images * 2. - 1.).latent_dist.sample() * 0.18215
         controlnet_images = rearrange(controlnet_images, "(b f) c h w -> b c f h w", f=num_controlnet_images)
 
@@ -517,7 +518,8 @@ with block as demo:
             gr.Markdown(instructions)      
 
 
-    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+    # device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+    device = torch.device("cuda")
     unet_path = 'models/unet.ckpt'
     image_controlnet_path = 'models/image_controlnet.ckpt'
     flow_controlnet_path = 'models/flow_controlnet.ckpt'