Update app.py

app.py

@@ -29,7 +29,7 @@ dtype = torch.float16
 pipe = AnimateDiffPipeline.from_pretrained(bases[base_loaded], torch_dtype=dtype).to(device)
 pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing", beta_schedule="linear")
 
-def generate_image(secret_token, prompt, base, motion, step):
+def generate_image(secret_token, prompt, base, width, height, motion, step):
     if secret_token != SECRET_TOKEN:
         raise gr.Error(
             f'Invalid secret token. Please fork the original space if you want to use it for yourself.')
@@ -57,22 +57,14 @@ def generate_image(secret_token, prompt, base, motion, step):
             pipe.set_adapters(["motion"], [0.7])
         motion_loaded = motion
 
-    progress((0, step))
-    def progress_callback(i, t, z):
-        progress((i+1, step))
-
     output = pipe(
         prompt=prompt,
 
-        # this corresponds roughly to 16:9
-        # which is the aspect ratio video used by AiTube
-        width=912, # 1024,
-        height=512, # 576,
+        width=width,
+        height=height,
         
         guidance_scale=1.0,
         num_inference_steps=step,
-        callback=progress_callback,
-        callback_steps=1
     )
 
     name = str(uuid.uuid4()).replace("-", "")
@@ -124,6 +116,20 @@ with gr.Blocks() as demo:
                 value=base_loaded,
                 interactive=True
             )
+            width = gr.Slider(
+                label='Width',
+                minimum=128,
+                maximum=2048,
+                step=32,
+                value=1024,
+            )
+            height = gr.Slider(
+                label='Height',
+                minimum=128,
+                maximum=2048,
+                step=32,
+                value=1024,
+            )
             select_motion = gr.Dropdown(
                 label='Motion',
                 choices=[
@@ -156,7 +162,7 @@ with gr.Blocks() as demo:
 
     submit.click(
         fn=generate_image,
-        inputs=[secret_token, prompt, select_base, select_motion, select_step],
+        inputs=[secret_token, prompt, select_base, width, height, select_motion, select_step],
         outputs=output_video_base64,
     )