Update app.py

app.py

@@ -6,7 +6,7 @@ from PIL import Image, ImageDraw
 from gradio_client import Client, handle_file
 import numpy as np
 import cv2
-
+import os
 
 # Инициализация моделей
 from transformers import OneFormerProcessor, OneFormerForUniversalSegmentation
@@ -21,7 +21,6 @@ oneFormer_model = OneFormerForUniversalSegmentation.from_pretrained("shi-labs/on
 # 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)
@@ -62,42 +61,6 @@ def segment_image(image):
 
     return cropped_masks_with_labels
 
-# def merge_segments_by_labels(gallery_images, labels_input):
-#     """
-#     Объединяет сегменты из галереи изображений в одно изображение,
-#     основываясь на введенных пользователем метках.
-
-#     Args:
-#         gallery_images: Список изображений сегментов (кортежи (изображение, метка)).
-#         labels_input: Строка с метками, разделенными точкой с запятой.
-
-#     Returns:
-#         Список изображений, где выбранные сегменты объединены в одно.
-#     """
-#     labels_to_merge = [label.strip() for label in labels_input.split(";")]
-#     merged_image = None
-#     merged_indices = []
-
-#     for i, (image_path, label) in enumerate(gallery_images):
-#         if label in labels_to_merge:
-#             image = cv2.imread(image_path)
-#             if merged_image is None:
-#                 merged_image = image.copy()
-#             else:
-#                 merged_image = cv2.add(merged_image, image)
-#             merged_indices.append(i)
-#     if merged_image is not None:
-#         new_gallery_images = [
-#             item for i, item in enumerate(gallery_images) if i not in merged_indices
-#         ]
-        
-#         new_name = labels_to_merge[0]
-#         new_gallery_images.append((merged_image, new_name))
-        
-#         return new_gallery_images
-#     else:
-#         return gallery_images
-
 
 def merge_segments_by_labels(gallery_images, labels_input):
     labels_to_merge = [label.strip() for label in labels_input.split(";")]
@@ -107,19 +70,19 @@ def merge_segments_by_labels(gallery_images, labels_input):
     for i, (image_path, label) in enumerate(gallery_images):  # Исправлено: image_path
         if label in labels_to_merge:
             # Загружаем изображение с помощью PIL, сохраняя альфа-канал
-            image = Image.open(image_path).convert("RGBA")  
+            image = Image.open(image_path).convert("RGBA")
 
             if merged_image is None:
                 merged_image = image.copy()
             else:
                 # Объединяем изображения с учетом альфа-канала
-                merged_image = Image.alpha_composite(merged_image, image)  
+                merged_image = Image.alpha_composite(merged_image, image)
             merged_indices.append(i)
-    
+
     if merged_image is not None:
         # Преобразуем объединенное изображение в numpy array
-        merged_image_np = np.array(merged_image)  
-        
+        merged_image_np = np.array(merged_image)
+
         new_gallery_images = [
             item for i, item in enumerate(gallery_images) if i not in merged_indices
         ]
@@ -130,42 +93,26 @@ def merge_segments_by_labels(gallery_images, labels_input):
         return gallery_images
 
 
-# def set_client_for_session(request: gr.Request):
-#     x_ip_token = request.headers['x-ip-token']
-#     return Client("JeffreyXiang/TRELLIS", headers={"X-IP-Token": x_ip_token})
-
 def set_hunyuan_client(request: gr.Request):
     try:
         x_ip_token = request.headers['x-ip-token']
-        return Client("tencent/Hunyuan3D-2", headers={"X-IP-Token": x_ip_token})
+        client = Client("tencent/Hunyuan3D-2", headers={"X-IP-Token": x_ip_token})
+        print(x_ip_token, "tencent/Hunyuan3D-2 Ip token")
+        return client
     except:
+        print("tencent/Hunyuan3D-2 no token")
         return Client("tencent/Hunyuan3D-2")
-    
+
 def set_vFusion_client(request: gr.Request):
     try:
         x_ip_token = request.headers['x-ip-token']
-        return Client("facebook/VFusion3D", headers={"X-IP-Token": x_ip_token})
+        client = Client("facebook/VFusion3D", headers={"X-IP-Token": x_ip_token})
+        print(x_ip_token, "facebook/VFusion3D Ip token")
+        return client
     except:
+        print("facebook/VFusion3D no token")
         return Client("facebook/VFusion3D")
 
-# def generate_3d_model(client, segment_output, segment_name):
-#     for i, (image_path, label) in enumerate(segment_output):
-#         if label == segment_name:
-#             result = client.predict(
-#                 image=handle_file(image_path),
-#                 multiimages=[],
-#                 seed=0,
-#                 ss_guidance_strength=7.5,
-#                 ss_sampling_steps=12,
-#                 slat_guidance_strength=3,
-#                 slat_sampling_steps=12,
-#                 multiimage_algo="stochastic",
-#                 api_name="/image_to_3d"
-#             )
-#             break
-#     print(result)
-#     return result["video"]
-
 def generate_3d_model(client, segment_output, segment_name):
     for i, (image_path, label) in enumerate(segment_output):
         if label == segment_name:
@@ -180,7 +127,7 @@ def generate_3d_model(client, segment_output, segment_name):
                 api_name="/shape_generation"
             )
             print(result)
-            return (result[0], result[0])
+            return result[0]
 
 def generate_3d_model_texture(client, segment_output, segment_name):
     for i, (image_path, label) in enumerate(segment_output):
@@ -196,7 +143,7 @@ def generate_3d_model_texture(client, segment_output, segment_name):
                 api_name="/generation_all"
             )
             print(result)
-            return (result[1], result[1])
+            return result[1]
 
 def generate_3d_model2(client, segment_output, segment_name):
     for i, (image_path, label) in enumerate(segment_output):
@@ -206,7 +153,7 @@ def generate_3d_model2(client, segment_output, segment_name):
                 api_name="/step_1_generate_obj"
             )
             print(result)
-            return (result[0], result[0])
+            return result[0]
 
 
 # def classify_segments(segments):
@@ -225,8 +172,6 @@ def generate_3d_model2(client, segment_output, segment_name):
 #     return inpainted
 
 
-from gradio_litmodel3d import LitModel3D
-
 with gr.Blocks() as demo:
     hunyuan_client = gr.State()
     vFusion_client = gr.State()
@@ -264,11 +209,10 @@ with gr.Blocks() as demo:
                 #     interactive=True     # this allow users to interact with the model
                 # )
                 trellis_output = gr.Model3D(label="3D Model")
-                trellis_output2 = gr.Model3D(clear_color=[0.0, 0.0, 0.0, 0.0],  label="3D Model Wireframe")
                 # trellis_button.click(generate_3d_model, inputs=[client, segment_output, trellis_input], outputs=trellis_output)
-                hunyuan_button.click(generate_3d_model, inputs=[hunyuan_client, segment_output, trellis_input], outputs=[trellis_output, trellis_output2])
-                hunyuan_button_texture.click(generate_3d_model_texture, inputs=[hunyuan_client, segment_output, trellis_input], outputs=[trellis_output, trellis_output2])
-                vFusion_button.click(generate_3d_model2, inputs=[vFusion_client, segment_output, trellis_input], outputs=[trellis_output, trellis_output2])
+                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)
 
         segment_button.click(segment_image, inputs=image_input, outputs=segment_output)
         # segment_button.click(segment_full_image, inputs=image_input, outputs=segment_output)