Update app.py

app.py

@@ -15,7 +15,8 @@ dashscope.api_key = DASHSCOPE_API_KEY
 KEEP_SUCCESS_TASK = 3600 * 10
 KEEP_RUNING_TASK = 3600 * 1
 # the total running task number in 1800 seconds
-LIMIT_RUNING_TASK = 4
+LIMIT_RUNING_TASK = 2
+LIMIT_HISTORY_RUNING_TASK = 4
 
 task_status = {}
 
@@ -32,7 +33,7 @@ def get_submit_code():
     submit_code = random.randint(0, 2147483647)
     #
     for sub_c, sub_info in copy.deepcopy(total_task_info["latest_1h_submit_status"]).items():
-        if time.time() - sub_info["time"] > 3600:
+        if time.time() - sub_info > 3600:
             total_task_info["latest_1h_submit_status"].pop(sub_c)
     return submit_code
 def t2v_generation(prompt, resolution, watermark_wan, seed=-1):
@@ -40,7 +41,7 @@ def t2v_generation(prompt, resolution, watermark_wan, seed=-1):
     total_task_info["latest_1h_submit_status"][get_submit_code()] = time.time()
     total_task_info["total_submit"] += 1
     if not allow_task_num():
-        gr.Info(f"Warning: The number of running tasks is too large, please wait for a while.")
+        gr.Info(f"Warning: The number of running tasks is too large, the estimate waiting time is {get_waiting_time('-1')} s.")
         return None, gr.Button(visible=True)
     try:
         rsp = VideoSynthesis.call(model="wanx2.1-t2v-plus", prompt=prompt, seed=seed,
@@ -55,11 +56,12 @@ def t2v_generation(prompt, resolution, watermark_wan, seed=-1):
 def t2v_generation_async(prompt, size, watermark_wan, seed=-1):
     print(seed)
     seed = seed if seed >= 0 else random.randint(0, 2147483647)
+    total_task_info["latest_1h_submit_status"][get_submit_code()] = time.time()
     total_task_info["total_submit"] += 1
     print(seed)
     if not allow_task_num():
-        gr.Info(f"Warning: The number of running tasks is too large, the estimate waiting time is {waiting_time('-1')} s.")
-        return None, False, gr.Button(visible=True), gr.Button(visible=False)
+        gr.Info(f"Warning: The number of running tasks is too large, the estimate waiting time is {get_waiting_time('-1')} s.")
+        return None, False, gr.Button(visible=True), gr.Button(visible=False), gr.Slider(), gr.Slider()
     try:
         rsp = VideoSynthesis.async_call(model="wanx2.1-t2v-plus",
                                         prompt=prompt,
@@ -68,10 +70,10 @@ def t2v_generation_async(prompt, size, watermark_wan, seed=-1):
                                         watermark_wanx=watermark_wan)
         task_id = rsp.output.task_id
         status = False
-        return task_id, status, gr.Button(visible=False), gr.Button(visible=True)
+        return task_id, status, gr.Button(visible=False), gr.Button(visible=True), get_cost_time(task_id), get_waiting_time(task_id)
     except Exception as e:
         gr.Warning(f"Warning: {e}")
-        return None, True, gr.Button(), gr.Button()
+        return None, True, gr.Button(), gr.Button(), gr.Slider(), gr.Slider()
 
 
 def i2v_generation(prompt, image, watermark_wan, seed=-1):
@@ -90,21 +92,21 @@ def i2v_generation(prompt, image, watermark_wan, seed=-1):
 
 def i2v_generation_async(prompt, image, watermark_wan, seed=-1):
     seed = seed if seed >= 0 else random.randint(0, 2147483647)
+    total_task_info["latest_1h_submit_status"][get_submit_code()] = time.time()
     total_task_info["total_submit"] += 1
     if not allow_task_num():
-        gr.Info(f"Warning: The number of running tasks is too large, the estimate waiting time is {waiting_time('-1')} s.")
-        return "", None, gr.Button(visible=True), gr.Button(visible=False)
+        gr.Info(f"Warning: The number of running tasks is too large, the estimate waiting time is {get_waiting_time('-1')} s.")
+        return "", None, gr.Button(visible=True), gr.Button(visible=False), gr.Slider(), gr.Slider()
     try:
         rsp = VideoSynthesis.async_call(model="wanx2.1-i2v-plus", prompt=prompt, seed=seed,
                                         img_url=image, watermark_wanx=watermark_wan)
         print(rsp)
         task_id = rsp.output.task_id
         status = False
-        return task_id, status, gr.Button(visible=False), gr.Button(visible=True)
+        return task_id, status, gr.Button(visible=False), gr.Button(visible=True), gr.Button(visible=True), get_cost_time(task_id), get_waiting_time(task_id)
     except Exception as e:
         gr.Warning(f"Warning: {e}")
-        return "", None, gr.Button(), gr.Button()
-
+        return "", None, gr.Button(), gr.Button(), gr.Slider(), gr.Slider()
 
 def get_result_with_task_id(task_id):
     if task_id == "": return True, None
@@ -122,7 +124,7 @@ def get_result_with_task_id(task_id):
         if status:
             total_task_info["total_complete_task"] += 1
             total_task_info["total_process_cost"] += time.time() - task_status[task_id]["time"]
-            print(total_task_info)
+            print(total_task_info["total_complete_task"], total_task_info["total_process_cost"])
     except:
         video_url = None
         status = False
@@ -139,7 +141,7 @@ def allow_task_num():
             num += 1
         if not task_status[task_id]["status"]:
             total_num += 1
-    return num < LIMIT_RUNING_TASK or num < 10
+    return num < LIMIT_RUNING_TASK or total_num < LIMIT_HISTORY_RUNING_TASK
 
 def get_waiting_time(task_id):
     # if the num of running task < Limit
@@ -154,9 +156,13 @@ def get_waiting_time(task_id):
     latest_submit_tasks = len(total_task_info["latest_1h_submit_status"])
     print("latest submit tasks", latest_submit_tasks)
     if task_id in task_status:
-        return int(num * 480 - (time.time() - task_status[task_id]["time"]))
+        return int(640 - (time.time() - task_status[task_id]["time"]))
     else:
-        return int(latest_submit_tasks * (total_task_info["total_process_cost"]/total_task_info["total_complete_task"]))
+        return int(latest_submit_tasks * (total_task_info["total_process_cost"]/(total_task_info["total_complete_task"]+1)))
+
+def online_get_waiting_time(task, t2v_task_id, i2v_task_id):
+    task_id = t2v_task_id if task == "t2v" else i2v_task_id
+    return get_waiting_time(task_id)
 
 def clean_task_status():
     # clean the task over 1800 seconds
@@ -174,13 +180,18 @@ def clean_task_status():
 
 def get_cost_time(task_id):
     if task_id in task_status and not task_status[task_id]["status"]:
-        et = time.time() - task_status[task_id]["time"]
+        et = int(time.time() - task_status[task_id]["time"])
         return f"{et:.2f}"
     else:
         return gr.Textbox()
 
+def online_get_cost_time(task, t2v_task_id, i2v_task_id):
+    task_id = t2v_task_id if task == "t2v" else i2v_task_id
+    return get_cost_time(task_id)
 
-def get_process_bar(task_id, status):
+
+def get_process_bar(task, t2v_task_id, i2v_task_id, status):
+    task_id = t2v_task_id if task == "t2v" else i2v_task_id
     clean_task_status()
     if task_id not in task_status:
         task_status[task_id] = {
@@ -200,7 +211,7 @@ def get_process_bar(task_id, status):
     if task_status[task_id]["status"]:
         task_status[task_id]["value"] = 100
     else:
-        task_status[task_id]["value"] += 1
+        task_status[task_id]["value"] += 5
     if task_status[task_id]["value"] >= 100 and not task_status[task_id]["status"]:
         task_status[task_id]["value"] = 95
     # print(task_id, task_status[task_id], task_status)
@@ -220,7 +231,8 @@ with gr.Blocks() as demo:
                </div>
 
                """)
-    task_id = gr.State(value="")
+    t2v_task_id = gr.State(value="")
+    i2v_task_id = gr.State(value="")
     status = gr.State(value=False)
     task = gr.State(value="t2v")
     with gr.Row():
@@ -270,12 +282,12 @@ with gr.Blocks() as demo:
             with gr.Row():
                 watermark_wan = gr.Checkbox(label="Watermark", value=True, visible=True, container=False)
                 seed = gr.Number(label="Seed", value=-1, container=True)
-                cost_time = gr.Number(label="Cost Time(secs)", value=get_cost_time, interactive=False,
-                                      every=120, inputs=[task_id], container=True)
-                waiting_time = gr.Number(label="Estimated Waiting Time(secs)", value=get_waiting_time, interactive=False,
-                                      every=120, inputs=[task_id], container=True)
+                cost_time = gr.Number(label="Cost Time(secs)", value=online_get_cost_time, interactive=False,
+                                      every=120, inputs=[task, t2v_task_id, i2v_task_id], container=True)
+                waiting_time = gr.Number(label="Estimated Waiting Time(secs)", value=online_get_waiting_time, interactive=False,
+                                      every=120, inputs=[task, t2v_task_id, i2v_task_id], container=True)
                 process_bar = gr.Slider(show_label=True, label="", value=get_process_bar, maximum=100,
-                                        interactive=True, every=120, inputs=[task_id, status], container=True)
+                                        interactive=True, every=120, inputs=[task, t2v_task_id, i2v_task_id, status], container=True)
 
     fake_video = gr.Video(label='Examples', visible=False, interactive=False)
     with gr.Row(visible=True) as t2v_eg:
@@ -300,8 +312,11 @@ with gr.Blocks() as demo:
             return gr.Video(value=video_url), ret_t2v_btn, ret_i2v_btn, ret_t2v_status_btn, ret_i2v_status_btn
         return gr.Video(value=None), gr.Button(), gr.Button(), gr.Button(), gr.Button()
 
+    def online_process_change(task, t2v_task_id, i2v_task_id):
+        task_id = t2v_task_id if task == 't2v' else i2v_task_id
+        return process_change(task_id, task)
 
-    process_bar.change(process_change, inputs=[task_id, task],
+    process_bar.change(online_process_change, inputs=[task, t2v_task_id, i2v_task_id],
                        outputs=[result_gallery, run_t2v_button, run_i2v_button,
                                 t2v_refresh_status, i2v_refresh_status])
 
@@ -320,26 +335,34 @@ with gr.Blocks() as demo:
     run_t2v_button.click(
         fn=t2v_generation_async,
         inputs=[txt2vid_prompt, resolution, watermark_wan, seed],
-        outputs=[task_id, status, run_t2v_button, t2v_refresh_status],
+        outputs=[t2v_task_id, status, run_t2v_button, t2v_refresh_status],
     )
 
     run_i2v_button.click(
         fn=i2v_generation_async,
         inputs=[img2vid_prompt, img2vid_image, watermark_wan, seed],
-        outputs=[task_id, status, run_i2v_button, i2v_refresh_status],
+        outputs=[i2v_task_id, status, run_i2v_button, i2v_refresh_status],
     )
 
     def status_refresh(task_id, task, status):
-        cost_time = get_cost_time(task_id)
+        if task_id in task_status and not task_status[task_id]["status"]:
+            cost_time = int(time.time() - task_status[task_id]["time"])
+        else:
+            cost_time = 0
+        status, video_url = get_result_with_task_id(task_id)
+        task_status[task_id]["status"] = status
+        task_status[task_id]["url"] = video_url
         waiting_time = get_waiting_time(task_id)
-        process_bar = get_process_bar(task_id, status)
+        value = task_status.get(task_id, {"value": 100})["value"]
+        value = min(value, int(cost_time*100/waiting_time))
+        process_bar = gr.Slider(label=f"({value}%)Generating" if value % 2 == 1 else f"({value}%)Generating.....", value=value)
         process_change_ret = process_change(task_id, task)
-        return process_change_ret + [cost_time, waiting_time, process_bar]
+        return *process_change_ret, cost_time, waiting_time, process_bar
 
 
     t2v_refresh_status.click(
         fn=status_refresh,
-        inputs=[task_id, task, status],
+        inputs=[t2v_task_id, task, status],
         outputs=[result_gallery, run_t2v_button, run_i2v_button,
                 t2v_refresh_status, i2v_refresh_status,
                 cost_time, waiting_time, process_bar]
@@ -347,7 +370,7 @@ with gr.Blocks() as demo:
 
     i2v_refresh_status.click(
         fn=status_refresh,
-        inputs=[task_id, task, status],
+        inputs=[i2v_task_id, task, status],
         outputs=[result_gallery, run_t2v_button, run_i2v_button,
                  t2v_refresh_status, i2v_refresh_status,
                  cost_time, waiting_time, process_bar]