Update app.py

app.py

@@ -115,10 +115,13 @@ def generate(segment, image, prompt, size, guidance_scale, num_inference_steps,
         num_channels_latents = int(num_channels_latents / 2)
         image = Image.open(image).convert('RGB')
         image.resize((size,size), Image.LANCZOS)
-        with torch.no_grad():
-            image = pipe.video_processor.preprocess(image, height=size, width=size).to(
-                device, dtype=prompt_embeds.dtype
-            )
+        
+        #with torch.no_grad():
+        pipe.vae.to("cpu")
+        image = pipe.video_processor.preprocess(image, height=size, width=size).to(
+            device, dtype=prompt_embeds.dtype
+        )
+        pipe.vae.to("cuda")
         num_latent_frames = (frames - 1) // pipe.vae_scale_factor_temporal + 1
         latents = pipe.prepare_latents(
             batch_size=1, num_channels_latents=pipe.transformer.config.in_channels, height=size, width=size, num_frames=frames,