add some features

Rodin.py

@@ -1,6 +1,7 @@
 import socketio
 import requests
 import json
+import time
 import random
 import base64
 import io
@@ -135,8 +136,12 @@ class Generator:
             
             mesh_response = rodin_mesh(prompt=prompt, group_uuid=None, settings=settings, images=images, name="images.jpeg", token=self.token)
             progress_checker = JobStatusChecker(BASE_URL, mesh_response['job']['subscription_key'])
-            progress_checker.start()
-            
+            try:
+                progress_checker.start()
+            except Exception as e:
+                print(f"Error in generating mesh: {e}")
+                time.sleep(5)
+                            
             task_uuid = mesh_response['uuid'] # The task_uuid should be same during whole generation process
         else:
             new_prompt = prompt
@@ -151,7 +156,11 @@ class Generator:
             # Check progress
             subscription_key = update_response['job']['subscription_key']
             checker = JobStatusChecker(BASE_URL, subscription_key)
-            checker.start()
+            try:
+                checker.start()
+            except Exception as e:
+                print(f"Error in updating mesh: {e}")
+                time.sleep(5)
         
         preview_image = next(reversed(rodin_history(task_uuid, self.token).items()))[1]["preview_image"]
         # print(f"Preview image URL: {rodin_history(task_uuid, self.token)}")

app.py

@@ -91,8 +91,8 @@ def handle_prompt_change(prompt):
 
     return hint_list
 
-def clear_task_uuid():
-    return ""
+def clear_task():
+    return "", "", [], None
 
 def return_render(image):
     image = Image.fromarray(image)
@@ -155,14 +155,12 @@ with gr.Blocks() as demo:
     cache_task_uuid = gr.Text(value="", visible=False)
 
     # button_generate.click(fn=return_render, inputs=[block_image], outputs=[raw_image, fake3d])
-    block_image.change(
+    block_image.upload(
         fn=do_nothing, 
-        js=reset_button_name, 
+        js=change_button_name, 
         inputs=[cacha_empty],
         outputs=[cacha_empty]
-    ).then(fn=clear_task_uuid, outputs=[cache_task_uuid], show_progress="hidden")
-    
-    button_generate.click(
+    ).success(
         fn=generator.preprocess, 
         inputs=[block_prompt, block_image, cache_image_base64, cache_task_uuid], 
         outputs=[block_prompt, cache_image_base64],
@@ -171,12 +169,31 @@ with gr.Blocks() as demo:
         fn=generator.generate_mesh,
         inputs=[block_prompt, cache_image_base64, cache_task_uuid],
         outputs=[cache_raw_image, cache_task_uuid, fake3d],
-    ).success(
+    )
+    
+    block_image.clear(
+        fn=do_nothing, 
+        js=reset_button_name, 
+        inputs=[cacha_empty],
+        outputs=[cacha_empty]
+    ).then(fn=clear_task, outputs=[cache_task_uuid, block_prompt, block_prompt_hint, fake3d], show_progress="hidden")
+    
+    button_generate.click(
         fn=do_nothing, 
         js=change_button_name, 
         inputs=[cacha_empty],
         outputs=[cacha_empty]
+    ).success(
+        fn=generator.preprocess, 
+        inputs=[block_prompt, block_image, cache_image_base64, cache_task_uuid], 
+        outputs=[block_prompt, cache_image_base64],
+        show_progress="minimal"
+    ).success(
+        fn=generator.generate_mesh,
+        inputs=[block_prompt, cache_image_base64, cache_task_uuid],
+        outputs=[cache_raw_image, cache_task_uuid, fake3d],
     )
+    
     block_default.click(fn=crop_image_default, inputs=[cache_raw_image], outputs=fake3d, show_progress="minimal")
     block_metal.click(fn=crop_image_metal, inputs=[cache_raw_image], outputs=fake3d, show_progress="minimal")
     block_contrast.click(fn=crop_image_contrast, inputs=[cache_raw_image], outputs=fake3d, show_progress="minimal")