Update app.py

app.py

@@ -96,15 +96,19 @@ def resize_image_to_bucket(image: Union[Image.Image, np.ndarray], bucket_reso: T
     return image
 
 
-def generate_video(prompt: str, frame1: Image.Image, frame2: Image.Image, guidance_scale: float, num_frames: int, num_inference_steps: int) -> bytes:
+
+def generate_video(prompt: str, frame1: Image.Image, frame2: Image.Image, resolution: str, guidance_scale: float, num_frames: int, num_inference_steps: int) -> bytes:
     # Debugging print statements
     print(f"Frame 1 Type: {type(frame1)}")
     print(f"Frame 2 Type: {type(frame2)}")
+    print(f"Resolution: {resolution}")
+    
+    # Parse resolution
+    width, height = map(int, resolution.split('x'))
     
     # Load and preprocess frames
     cond_frame1 = np.array(frame1)
     cond_frame2 = np.array(frame2)
-    height, width = 720, 1280
     cond_frame1 = resize_image_to_bucket(cond_frame1, bucket_reso=(width, height))
     cond_frame2 = resize_image_to_bucket(cond_frame2, bucket_reso=(width, height))
     cond_video = np.zeros(shape=(num_frames, height, width, 3))
@@ -136,7 +140,6 @@ def generate_video(prompt: str, frame1: Image.Image, frame2: Image.Image, guidan
     with open(video_path, "rb") as video_file:
         video_bytes = video_file.read()
     return video_bytes
-
 @torch.inference_mode()
 def call_pipe(
     pipe,
@@ -301,11 +304,16 @@ def main():
         gr.Textbox(label="Prompt", value="a woman"),
         gr.Image(label="Frame 1", type="pil"),
         gr.Image(label="Frame 2", type="pil"),
+        gr.Dropdown(
+            label="Resolution",
+            choices=["720x1280", "544x960", "1280x720", "960x544", "720x720"],
+            value="544x960"
+        ),
         # gr.Textbox(label="Frame 1 URL", value="https://i-bacon.bunkr.ru/11b45aa7-630b-4189-996f-a6b37a697786.png"),
         # gr.Textbox(label="Frame 2 URL", value="https://i-bacon.bunkr.ru/2382224f-120e-482d-a75d-f1a1bf13038c.png"),
         gr.Slider(minimum=0.1, maximum=20, step=0.1, label="Guidance Scale", value=6.0),
-        gr.Slider(minimum=1, maximum=129, step=1, label="Number of Frames", value=77),
-        gr.Slider(minimum=1, maximum=100, step=1, label="Number of Inference Steps", value=20)
+        gr.Slider(minimum=1, maximum=129, step=1, label="Number of Frames", value=49),
+        gr.Slider(minimum=1, maximum=100, step=1, label="Number of Inference Steps", value=30)
     ]
     
     # Define the interface outputs