session based aspect ratio as oppose to the global one

app.py

@@ -321,8 +321,6 @@ 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]
     
@@ -330,7 +328,7 @@ class WanInferencePipeline(nn.Module):
             image = None
             select_size = [height, width]
         num = self.args.max_tokens * 16 * 16 * 4
-        den = select_size[0] * select_size[1]
+        den = h * w
         L0 = num // den 
         diff = (L0 - 1) % 4
         L  = L0 - diff
@@ -394,29 +392,22 @@ 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
-            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
+            h = height
+            w = width
 
         # step 1: numerator and denominator as ints
         num = args.max_tokens * 16 * 16 * 4
-        den = select_size[0] * select_size[1]
+        den = h * w
         
         # 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:
@@ -615,7 +606,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, session_id = None):
+def preprocess_img(input_image_path, raw_image_path, orientation_state, session_id = None):
 
     if session_id is None:
         session_id = uuid.uuid4().hex
@@ -631,7 +622,7 @@ def preprocess_img(input_image_path, raw_image_path, session_id = None):
     image = inferpipe.transform(image).unsqueeze(0).to(dtype=inferpipe.dtype)
     
     _, _, h, w = image.shape
-    select_size = match_size(getattr( args, f'image_sizes_{ args.max_hw}'), h, w)
+    select_size = match_size(orientation_state, h, w)
     image = resize_pad(image, (h, w), select_size)
     image = image * 2.0 - 1.0
     image = image[:, :, None]
@@ -649,13 +640,12 @@ def preprocess_img(input_image_path, raw_image_path, session_id = None):
 
 def infer_example(image_path, audio_path, text, num_steps, raw_image_path, session_id = None, progress=gr.Progress(track_tqdm=True),):
 
-    current_image_size = args.image_sizes_720
-    args.image_sizes_720 = [[720, 400]]
+    if session_id is None:
+        session_id = uuid.uuid4().hex
 
+    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)
@@ -713,7 +703,8 @@ def infer(image_path, audio_path, text, num_steps, session_id = None, progress=g
 
 def apply_image(request):
     print('image applied')
-    return request, None
+
+    return request, request
 
 def apply_audio(request):
     print('audio applied')
@@ -739,13 +730,15 @@ def orientation_changed(session_id, evt: gr.EventData):
     detail = getattr(evt, "data", None) or getattr(evt, "_data", {}) or {}
 
     if detail['value'] == "9:16":
-        args.image_sizes_720 = [[720, 400]]
+        orientation_state = [[720, 400]]
     elif detail['value'] == "1:1":
-        args.image_sizes_720 = [[720, 720]]
+        orientation_state = [[720, 720]]
     elif detail['value'] == "16:9":
-        args.image_sizes_720 = [[400, 720]]
+        orientation_state = [[400, 720]]
+
+    print(f'{session_id} has {orientation_state} orientation')
 
-    print(f'{session_id} has {args.image_sizes_720} orientation')
+    return orientation_state
 
 def clear_raw_image():
     return ''
@@ -819,6 +812,7 @@ 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])
 
 
@@ -936,7 +930,9 @@ with gr.Blocks(css=css) as demo:
                     ],
                     label="Image Samples",
                     inputs=[image_input],
-                    cache_examples=False
+                    outputs=[image_input, raw_img_text],
+                    fn=apply_image,
+                    cache_examples=True
                     )
 
                 audio_examples = gr.Examples(        
@@ -981,9 +977,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]).then(fn=preprocess_img, inputs=[image_input, raw_img_text, session_state], outputs=[image_input, raw_img_text])
+    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.clear(fn=clear_raw_image, outputs=[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.upload(fn=preprocess_img, inputs=[image_input, raw_img_text, orientation_state, 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])