Update app.py

app.py

@@ -20,7 +20,6 @@ oneFormer_model = OneFormerForUniversalSegmentation.from_pretrained("shi-labs/on
 # classification = pipeline("image-classification", model="google/vit-base-patch16-224")
 # upscaling_client = InferenceClient(model="stabilityai/stable-diffusion-x4-upscaler")
 # inpainting_client = InferenceClient(model="stabilityai/stable-diffusion-inpainting")
-
 # Функции для обработки изображений
 def segment_image(image):
     image = Image.fromarray(image)
@@ -93,7 +92,7 @@ def merge_segments_by_labels(gallery_images, labels_input):
         return gallery_images
 
 
-def set_hunyuan_client(request: gr.Request):
+def hunyuan_client(request: gr.Request):
     try:
         client = Client("tencent/Hunyuan3D-2", headers={"X-IP-Token": request.headers['x-ip-token']})
         print("tencent/Hunyuan3D-2 Ip token")
@@ -102,7 +101,7 @@ def set_hunyuan_client(request: gr.Request):
         print("tencent/Hunyuan3D-2 no token")
         return Client("tencent/Hunyuan3D-2")
 
-def set_vFusion_client(request: gr.Request):
+def vFusion_client(request: gr.Request):
     try:
         client = Client("facebook/VFusion3D", headers={"X-IP-Token": request.headers['x-ip-token']})
         print("facebook/VFusion3D Ip token")
@@ -111,9 +110,10 @@ def set_vFusion_client(request: gr.Request):
         print("facebook/VFusion3D no token")
         return Client("facebook/VFusion3D")
 
-def generate_3d_model(client, segment_output, segment_name):
+def generate_3d_model(segment_output, request: gr.Request):
     for i, (image_path, label) in enumerate(segment_output):
         if label == segment_name:
+            client = hunyuan_client(request)
             result = client.predict(
                 caption="",
                 image=handle_file(image_path),
@@ -127,11 +127,10 @@ def generate_3d_model(client, segment_output, segment_name):
             print(result)
             return result[0]
 
-def generate_3d_model_texture(client, segment_output, segment_name, request: gr.Request):
+def generate_3d_model_texture(segment_output, segment_name, request: gr.Request):
     for i, (image_path, label) in enumerate(segment_output):
         if label == segment_name:
-            client = Client("tencent/Hunyuan3D-2", headers={"X-IP-Token": request.headers['x-ip-token']})
-            print("tencent/Hunyuan3D-2 new client token")
+            client = hunyuan_client(request)
             result = client.predict(
                 caption="",
                 image=handle_file(image_path),
@@ -145,9 +144,10 @@ def generate_3d_model_texture(client, segment_output, segment_name, request: gr.
             print(result)
             return result[1]
 
-def generate_3d_model2(client, segment_output, segment_name):
+def generate_3d_model2(segment_output, segment_name, request: gr.Request):
     for i, (image_path, label) in enumerate(segment_output):
         if label == segment_name:
+            client = vFusion_client(request)
             result = client.predict(
                 image=handle_file(image_path),
                 api_name="/step_1_generate_obj"
@@ -172,10 +172,11 @@ def generate_3d_model2(client, segment_output, segment_name):
 #     return inpainted
 
 
-with gr.Blocks() as demo:
-    hunyuan_client = gr.State()
-    vFusion_client = gr.State()
 
+
+########## GRADIO ##########
+
+with gr.Blocks() as demo:
     gr.Markdown("# Анализ и редактирование помещений")
 
     with gr.Tab("Сканирование"):
@@ -197,9 +198,9 @@ with gr.Blocks() as demo:
             with gr.Column(scale=5):
                 trellis_output = gr.Model3D(label="3D Model")
                 
-                hunyuan_button.click(generate_3d_model, inputs=[hunyuan_client, segment_output, trellis_input], outputs=trellis_output)
-                hunyuan_button_texture.click(generate_3d_model_texture, inputs=[hunyuan_client, segment_output, trellis_input], outputs=trellis_output)
-                vFusion_button.click(generate_3d_model2, inputs=[vFusion_client, segment_output, trellis_input], outputs=trellis_output)
+                hunyuan_button.click(generate_3d_model, inputs=[segment_output, trellis_input], outputs=trellis_output)
+                hunyuan_button_texture.click(generate_3d_model_texture, inputs=[segment_output, trellis_input], outputs=trellis_output)
+                vFusion_button.click(generate_3d_model2, inputs=[segment_output, trellis_input], outputs=trellis_output)
 
         segment_button.click(segment_image, inputs=image_input, outputs=segment_output)
         # segment_button.click(segment_full_image, inputs=image_input, outputs=segment_output)
@@ -222,7 +223,4 @@ with gr.Blocks() as demo:
     #     model_button = gr.Button("Создать 3D модель")
     #     model_button.click(generate_3d_model, inputs=segment_input_3d, outputs=model_output)
 
-    demo.load(set_hunyuan_client, None, hunyuan_client)
-    demo.load(set_vFusion_client, None, vFusion_client)
-
-demo.launch(debug=True, show_error=True)
+demo.launch(debug=True, show_error=True)