Update app.py

app.py

@@ -321,6 +321,8 @@ class WanInferencePipeline(nn.Module):
             image = self.transform(image).unsqueeze(0).to(dtype=self.dtype)
             
             _, _, h, w = image.shape
+            select_size = match_size(getattr( self.args, f'image_sizes_{ self.args.max_hw}'), h, w)
+            image = resize_pad(image, (h, w), select_size)
             image = image * 2.0 - 1.0
             image = image[:, :, None]
     
@@ -328,7 +330,7 @@ class WanInferencePipeline(nn.Module):
             image = None
             select_size = [height, width]
         num = self.args.max_tokens * 16 * 16 * 4
-        den = h * w
+        den = select_size[0] * select_size[1]
         L0 = num // den 
         diff = (L0 - 1) % 4
         L  = L0 - diff
@@ -392,22 +394,29 @@ class WanInferencePipeline(nn.Module):
             image = self.transform(image).unsqueeze(0).to(self.device, dtype=self.dtype)
             
             _, _, h, w = image.shape
+            select_size = match_size(getattr(self.args, f'image_sizes_{self.args.max_hw}'), h, w)
+            image = resize_pad(image, (h, w), select_size)
             image = image * 2.0 - 1.0
             image = image[:, :, None]
 
         else:
             image = None
-            h = height
-            w = width
+            select_size = [height, width]
+        # L = int(self.args.max_tokens * 16 * 16 * 4 / select_size[0] / select_size[1])
+        # L = L // 4 * 4 + 1 if L % 4 != 0 else L - 3  # video frames
+        # T = (L + 3) // 4  # latent frames
 
         # step 1: numerator and denominator as ints
         num = args.max_tokens * 16 * 16 * 4
-        den = h * w
+        den = select_size[0] * select_size[1]
         
         # step 2: integer division
         L0 = num // den  # exact floor division, no float in sight
         
         # step 3: make it ≡ 1 mod 4
+        #    if L0 % 4 == 1, keep L0;
+        #    otherwise subtract the difference so that (L0 - diff) % 4 == 1,
+        #    but ensure the result stays positive.
         diff = (L0 - 1) % 4
         L  = L0 - diff
         if L < 1:
@@ -606,7 +615,7 @@ def get_duration(image_path, audio_path, text, num_steps, session_id, progress):
 
     return int(duration_s)
 
-def preprocess_img(input_image_path, raw_image_path, orientation_state, session_id = None):
+def preprocess_img(input_image_path, raw_image_path, session_id = None):
 
     if session_id is None:
         session_id = uuid.uuid4().hex
@@ -622,7 +631,7 @@ def preprocess_img(input_image_path, raw_image_path, orientation_state, session_
     image = inferpipe.transform(image).unsqueeze(0).to(dtype=inferpipe.dtype)
     
     _, _, h, w = image.shape
-    select_size = match_size(orientation_state, h, w)
+    select_size = match_size(getattr( args, f'image_sizes_{ args.max_hw}'), h, w)
     image = resize_pad(image, (h, w), select_size)
     image = image * 2.0 - 1.0
     image = image[:, :, None]
@@ -640,12 +649,13 @@ def preprocess_img(input_image_path, raw_image_path, orientation_state, session_
 
 def infer_example(image_path, audio_path, text, num_steps, raw_image_path, session_id = None, progress=gr.Progress(track_tqdm=True),):
 
-    if session_id is None:
-        session_id = uuid.uuid4().hex
+    current_image_size = args.image_sizes_720
+    args.image_sizes_720 = [[720, 400]]
 
-    image_path, _ = preprocess_img(image_path, image_path, [[720, 400]], session_id)
     result = infer(image_path, audio_path, text, num_steps, session_id, progress)
 
+    args.image_sizes_720 = current_image_size
+
     return result
     
 @spaces.GPU(duration=get_duration)
@@ -703,8 +713,7 @@ def infer(image_path, audio_path, text, num_steps, session_id = None, progress=g
 
 def apply_image(request):
     print('image applied')
-
-    return request, request
+    return request, None
 
 def apply_audio(request):
     print('audio applied')
@@ -730,15 +739,13 @@ def orientation_changed(session_id, evt: gr.EventData):
     detail = getattr(evt, "data", None) or getattr(evt, "_data", {}) or {}
 
     if detail['value'] == "9:16":
-        orientation_state = [[720, 400]]
+        args.image_sizes_720 = [[720, 400]]
     elif detail['value'] == "1:1":
-        orientation_state = [[720, 720]]
+        args.image_sizes_720 = [[720, 720]]
     elif detail['value'] == "16:9":
-        orientation_state = [[400, 720]]
-
-    print(f'{session_id} has {orientation_state} orientation')
+        args.image_sizes_720 = [[400, 720]]
 
-    return orientation_state
+    print(f'{session_id} has {args.image_sizes_720} orientation')
 
 def clear_raw_image():
     return ''
@@ -812,7 +819,6 @@ css = """
 with gr.Blocks(css=css) as demo:
 
     session_state = gr.State()
-    orientation_state = gr.State([[720, 400]])
     demo.load(start_session, outputs=[session_state])
 
 
@@ -930,9 +936,7 @@ with gr.Blocks(css=css) as demo:
                     ],
                     label="Image Samples",
                     inputs=[image_input],
-                    outputs=[image_input, raw_img_text],
-                    fn=apply_image,
-                    cache_examples=True
+                    cache_examples=False
                     )
 
                 audio_examples = gr.Examples(        
@@ -977,9 +981,9 @@ with gr.Blocks(css=css) as demo:
         inputs=[audio_input, limit_on, session_state],
         outputs=[audio_input],
     )
-    image_input.orientation(fn=orientation_changed, inputs=[session_state], outputs=[orientation_state]).then(fn=preprocess_img, inputs=[image_input, raw_img_text, orientation_state, session_state], outputs=[image_input, raw_img_text])
+    image_input.orientation(fn=orientation_changed, inputs=[session_state]).then(fn=preprocess_img, inputs=[image_input, raw_img_text, session_state], outputs=[image_input, raw_img_text])
     image_input.clear(fn=clear_raw_image, outputs=[raw_img_text])
-    image_input.upload(fn=preprocess_img, inputs=[image_input, raw_img_text, orientation_state, session_state], outputs=[image_input, raw_img_text])
+    image_input.upload(fn=preprocess_img, inputs=[image_input, raw_img_text, session_state], outputs=[image_input, raw_img_text])
     image_input.change(fn=update_generate_button, inputs=[image_input, audio_input, text_input, num_steps, session_state], outputs=[time_required])
     audio_input.change(fn=update_generate_button, inputs=[image_input, audio_input, text_input, num_steps, session_state], outputs=[time_required])
     num_steps.change(fn=slider_value_change, inputs=[image_input, audio_input, text_input, num_steps, session_state, adaptive_text], outputs=[time_required, text_input])