fixed gradio

app.py

@@ -41,7 +41,7 @@ def forward_diffusion_sample(x_0, t):
     return sqrt_alphas_cumprod_t.to(device) * x_0.to(device) \
     + sqrt_one_minus_alphas_cumprod_t.to(device) * noise.to(device), noise.to(device)
 
-T = 1000
+T = 10
 betas = cosine_beta_schedule(timesteps=T)
 # Pre-calculate different terms for closed form
 alphas = 1. - betas
@@ -89,7 +89,7 @@ torch.cuda.empty_cache()
 
 def save_video_frames_as_mp4(frames, fps, save_path):
     frame_h, frame_w = frames[0].shape[2:]
-    fourcc = cv2.VideoWriter_fourcc('m', 'p', '4', 'v')
+    fourcc = cv2.VideoWriter_fourcc(*'avc1')
     video = cv2.VideoWriter(save_path, fourcc, fps, (frame_w, frame_h))
     frames = frames[0]
     for frame in frames:
@@ -157,7 +157,8 @@ def get_image_embedding(input_image):
     return encoder_hidden_states
 
 def predict_fn(img_path, progress=gr.Progress()):
-    image = get_transform(Image.open(img_path).convert('RGB'))
+    img2tensor = get_transform()
+    image = img2tensor(img_path).unsqueeze(0).to(device)
     encoder_hidden_states = get_image_embedding(input_image=image)
     encoded_image = VAE_encode(image)
     noise_video = torch.randn([1, frameLimit, 4, 80, 64]).to(device)
@@ -171,19 +172,26 @@ def predict_fn(img_path, progress=gr.Progress()):
         save_video_frames_as_mp4(final_video, 25, "result.mp4")
     return "result.mp4"
 
-with gr.Tab("Image-to-Video"):
-    with gr.Row():
-        with gr.Column():
-            image_input = gr.Image(type="pil", label="Input Image")
-            img_generate = gr.Button("Generate Video")
-        with gr.Column():
-            img_output = gr.Video(label="Generated Video")
-    gr.Examples(
-        examples=[
-            ['sample/blue.jpg',]
-        ],
-        inputs=[image_input],
-        outputs=[img_output],
-        fn=predict_fn,
-        cache_examples='lazy',
-    )
+with gr.Blocks() as demo:
+    with gr.Tab("Image-to-Video"):
+        with gr.Row():
+            with gr.Column():
+                image_input = gr.Image(type="pil", label="Input Image")
+                img_generate = gr.Button("Generate Video")
+            with gr.Column():
+                img_output = gr.Video(label="Generated Video")
+        gr.Examples(
+            examples=[
+                ['sample/blue.jpg',]
+            ],
+            inputs=[image_input],
+            outputs=[]
+        )
+
+        img_generate.click(
+            fn=predict_fn,
+            inputs=image_input,
+            outputs=img_output
+        )
+
+demo.launch()