Update app.py

app.py

@@ -15,12 +15,33 @@ 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
 
-def t2v_generation(prompt, resolution, watermark_wan, seed = -1):
+task_status = {}
+
+total_task_info = {
+    "total_process_cost": 0,
+    "total_complete_task": 0,
+    "total_submit": 0,
+    "latest_1h_submit_status": {
+
+    }
+}
+
+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 > 3600:
+            total_task_info["latest_1h_submit_status"].pop(sub_c)
+    return submit_code
+def t2v_generation(prompt, resolution, 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, 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,
@@ -31,13 +52,16 @@ def t2v_generation(prompt, resolution, watermark_wan, seed = -1):
         gr.Warning(f"Warning: {e}")
         return None, gr.Button(visible=True)
 
-def t2v_generation_async(prompt, size, 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, please wait for a while.")
-        return None, False, gr.Button(visible=True)
+        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,
@@ -46,17 +70,18 @@ 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)
+        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()
+        return None, True, gr.Button(), gr.Button(), gr.Slider(), gr.Slider()
 
-def i2v_generation(prompt, image, watermark_wan, seed = -1):
+
+def i2v_generation(prompt, image, watermark_wan, seed=-1):
     seed = seed if seed >= 0 else random.randint(0, 2147483647)
     video_url = None
     try:
-        rsp = VideoSynthesis.call(model="wanx2.1-i2v-plus", prompt=prompt, img_url= image,
-                                  seed = seed,
+        rsp = VideoSynthesis.call(model="wanx2.1-i2v-plus", prompt=prompt, img_url=image,
+                                  seed=seed,
                                   watermark_wanx=watermark_wan
                                   )
         video_url = rsp.output.video_url
@@ -64,26 +89,29 @@ def i2v_generation(prompt, image, watermark_wan, seed = -1):
         gr.Warning(f"Warning: {e}")
     return video_url
 
-def i2v_generation_async(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, please wait for a while.")
-        return "", None, gr.Button(visible=True)
+        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)
+                                        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)
+        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()
+        return "", None, gr.Button(), gr.Button(), gr.Slider(), gr.Slider()
 
 def get_result_with_task_id(task_id):
     if task_id == "": return True, None
     try:
-        rsp = VideoSynthesis.fetch(task = task_id)
+        rsp = VideoSynthesis.fetch(task=task_id)
         print(rsp)
         if rsp.output.task_status == "FAILED":
             gr.Info(f"Warning: task running {rsp.output.task_status}")
@@ -93,12 +121,17 @@ def get_result_with_task_id(task_id):
             video_url = rsp.output.video_url
             video_url = video_url if video_url != "" else None
             status = video_url is not None
+        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["total_complete_task"], total_task_info["total_process_cost"])
     except:
         video_url = None
         status = False
-    return status, None if video_url=="" else video_url
+    return status, None if video_url == "" else video_url
     # return True, "https://dashscope-result-wlcb.oss-cn-wulanchabu.aliyuncs.com/1d/f8/20250220/e7d3f375/ccc590a2-7e90-4d92-84bc-22668db42979.mp4?Expires=1740137152&OSSAccessKeyId=LTAI5tQZd8AEcZX6KZV4G8qL&Signature=i3S3jA5FY6XYfvzZNHnvQiPzZSw%3D"
-task_status = {}
+
+
 
 def allow_task_num():
     num = 0
@@ -108,7 +141,25 @@ 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
+    # waiting time = num * 480s
+    # task_id not in task_status, return a large number
+    # prediction the waiting time
+    # avg_cost * latest submit time
+    num = 0
+    for task_id in task_status:
+        if not task_status[task_id]["status"]:
+            num += 1
+    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(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"]+1)))
+
 def clean_task_status():
     # clean the task over 1800 seconds
     for task_id in copy.deepcopy(task_status):
@@ -122,9 +173,10 @@ def clean_task_status():
             if task_status[task_id]["time"] + KEEP_RUNING_TASK < time.time():
                 task_status.pop(task_id)
 
-def cost_time(task_id):
+
+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()
@@ -141,7 +193,7 @@ def get_process_bar(task_id, status):
         }
     if not task_status[task_id]["status"]:
         # only when > 50% do check status
-        if task_status[task_id]["value"] >= 10 and task_status[task_id]["value"] % 5 == 0:
+        if task_status[task_id]["value"] >= 5 and task_status[task_id]["value"] % 5 == 0:
             status, video_url = get_result_with_task_id(task_id)
         else:
             status, video_url = False, None
@@ -150,16 +202,15 @@ 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)
     value = task_status[task_id]["value"]
-    return gr.Slider(label= f"({value}%)Generating" if value%2==1 else f"({value}%)Generating.....", value=value)
+    return gr.Slider(label=f"({value}%)Generating" if value % 2 == 1 else f"({value}%)Generating.....", value=value)
 
 
 with gr.Blocks() as demo:
-    
     gr.HTML("""
                <div style="text-align: center; font-size: 32px; font-weight: bold; margin-bottom: 20px;">
                    Wan2.1: Open and Advanced Large-Scale Video Generative Models
@@ -169,7 +220,7 @@ with gr.Blocks() as demo:
                    <a href="https://huggingface.co/Wan-AI">Huggingface</a> |
                    <a href="https://modelscope.cn/organization/Wan-AI">Modelscope</a>
                </div>
-               
+
                """)
     task_id = gr.State(value="")
     status = gr.State(value=False)
@@ -194,6 +245,7 @@ with gr.Blocks() as demo:
                             )
                         with gr.Row():
                             run_t2v_button = gr.Button("Generate Video")
+                            t2v_refresh_status = gr.Button("Refresh Generating Status", visible=False)
                     # Image to Video Tab
                     with gr.TabItem("Image to Video") as i2v_tab:
                         with gr.Row():
@@ -211,18 +263,21 @@ with gr.Blocks() as demo:
                                 )
                         with gr.Row():
                             run_i2v_button = gr.Button("Generate Video")
+                            i2v_refresh_status = gr.Button("Refresh Generating Status", visible=False)
         with gr.Column():
             with gr.Row():
                 result_gallery = gr.Video(label='Generated Video',
                                           interactive=False,
                                           height=500)
             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=cost_time, interactive=False,
-                                           every=10, inputs=[task_id], container=True)
-                    process_bar = gr.Slider(show_label=True, label="", value=get_process_bar, maximum=100,
-                      interactive=True, every=10, inputs=[task_id, status], container=True)
+                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)
+                process_bar = gr.Slider(show_label=True, label="", value=get_process_bar, maximum=100,
+                                        interactive=True, every=120, inputs=[task_id, status], container=True)
 
     fake_video = gr.Video(label='Examples', visible=False, interactive=False)
     with gr.Row(visible=True) as t2v_eg:
@@ -235,38 +290,77 @@ with gr.Blocks() as demo:
                     inputs=[img2vid_prompt, img2vid_image, result_gallery],
                     outputs=[result_gallery])
 
+
     def process_change(task_id, task):
         status = task_status[task_id]["status"]
         if status:
             video_url = task_status[task_id]["url"]
             ret_t2v_btn = gr.Button(visible=True) if task == 't2v' else gr.Button()
+            ret_t2v_status_btn = gr.Button(visible=False) if task == 't2v' else gr.Button()
             ret_i2v_btn = gr.Button(visible=True) if task == 'i2v' else gr.Button()
-            return gr.Video(value=video_url), ret_t2v_btn, ret_i2v_btn
-        return gr.Video(value=None), gr.Button(), gr.Button()
+            ret_i2v_status_btn = gr.Button(visible=False) if task == 'i2v' else gr.Button()
+            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()
+
 
     process_bar.change(process_change, inputs=[task_id, task],
-                       outputs=[result_gallery, run_t2v_button, run_i2v_button])
+                       outputs=[result_gallery, run_t2v_button, run_i2v_button,
+                                t2v_refresh_status, i2v_refresh_status])
 
 
     def switch_i2v_tab():
         return gr.Row(visible=False), gr.Row(visible=True), "i2v"
 
+
     def switch_t2v_tab():
         return gr.Row(visible=True), gr.Row(visible=False), "t2v"
 
+
     i2v_tab.select(switch_i2v_tab, outputs=[t2v_eg, i2v_eg, task])
     t2v_tab.select(switch_t2v_tab, outputs=[t2v_eg, i2v_eg, task])
 
     run_t2v_button.click(
         fn=t2v_generation_async,
         inputs=[txt2vid_prompt, resolution, watermark_wan, seed],
-        outputs=[task_id, status, run_t2v_button],
+        outputs=[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],
+        outputs=[task_id, status, run_i2v_button, i2v_refresh_status],
+    )
+
+    def status_refresh(task_id, task, status):
+        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)
+        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
+
+
+    t2v_refresh_status.click(
+        fn=status_refresh,
+        inputs=[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]
+    )
+
+    i2v_refresh_status.click(
+        fn=status_refresh,
+        inputs=[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]
     )
 
 demo.queue(max_size=10)