Update main.py

main.py

@@ -33,7 +33,7 @@ url_pre = "https://ap-east-1.tensorart.cloud/v1"
 SAVE_DIR = "generated_images"
 Path(SAVE_DIR).mkdir(exist_ok=True)
 
-# Danh sách sản phẩm (copy từ mã gốc)
+# Danh sách sản phẩm
 PRODUCT_GROUPS = {
     "Standard": {
         "C1012 Glacier White": "817687427545199895",
@@ -48,30 +48,8 @@ PRODUCT_GROUPS = {
         "C3105 Casla Cloudy": "828912225788997963",
         "C3146 Casla Nova": "828013009961087650",
         "C2240 Marquin": "828085015087780649",
-        "C2262 Concrete (Honed)": "822211862058871636",
-        "C3311 Calacatta Sky": "829984593223502930",
-        "C3346 Massimo": "827938741386607132",
-    },
-    "Luxury": {
-        "C4143 Mario": "829984593223502930",
-        "C4145 Marina": "828132560375742058",
-        "C4202 Calacatta Gold": "828167757632695310",
-        "C1205 Casla Everest": "828296778450463190",
-        "C4211 Calacatta Supreme": "828436321937882328",
-        "C4204 Calacatta Classic": "828422973179466146",
-        "C5240 Spring": "is coming",
-        "C1102 Super White": "828545723344775887",
-        "C4246 Casla Mystery": "828544778451950698",
-        "C4345 Oro": "828891068780182635",
-        "C4346 Luxe": "829436426547535131",
-        "C4342 Casla Eternal": "829190256201829181",
-        "C4221 Athena": "829644354504131520",
-        "C4222 Lagoon": "is coming",
-        "C5225 Amber": "is coming",
-    },
-    "Super Luxury": {
-        "C4255 Calacatta Extra": "829659013227537217",
     },
+    # Thêm các nhóm khác nếu cần
 }
 
 PRODUCT_IMAGE_MAP = {
@@ -199,125 +177,31 @@ def txt2img(prompt, width, height, product_codes):
 def generate_mask(image_resource_id, position, selected_product_code):
     try:
         if not image_resource_id:
-            raise Exception("Không có image_resource_id hợp lệ - ảnh gốc chưa được upload")
-        print(f"Using image_resource_id: {image_resource_id}")
-        time.sleep(10)  # Đợi đồng bộ tài nguyên
-        
+            raise Exception("Invalid image_resource_id")
+        time.sleep(10)
         short_code = selected_product_code.split()[0]
         texture_filepath = PRODUCT_IMAGE_MAP.get(selected_product_code)
-        print(f"Texture file: {texture_filepath}, exists: {os.path.exists(texture_filepath)}")
         if not texture_filepath or not os.path.exists(texture_filepath):
-            raise Exception(f"Không tìm thấy ảnh sản phẩm cho mã {short_code}")
+            raise Exception(f"Texture file not found for {short_code}")
         
         texture_resource_id = upload_image_to_tensorart(texture_filepath)
-        print(f"Texture resource_id: {texture_resource_id}")
         if not texture_resource_id:
-            raise Exception(f"Không thể upload ảnh sản phẩm {short_code}")
-        time.sleep(10)  # Đợi đồng bộ tài nguyên
-
-        if isinstance(position, (set, list)):
-            position = position[0] if position else "default"
-        print(f"Position: {position}, type: {type(position)}")
+            raise Exception(f"Failed to upload texture for {short_code}")
+        time.sleep(10)
 
-        # Dùng params đúng như mẫu TensorArt
         workflow_params = {
-            "1": {
-                "classType": "LayerMask: SegmentAnythingUltra V3",
-                "inputs": {
-                    "black_point": 0.3,
-                    "detail_dilate": 6,
-                    "detail_erode": 65,
-                    "detail_method": "GuidedFilter",
-                    "device": "cuda",
-                    "image": ["2", 0],
-                    "max_megapixels": 2,
-                    "process_detail": True,
-                    "prompt": ["4", 0],
-                    "sam_models": ["3", 0],
-                    "threshold": 0.3,
-                    "white_point": 0.99
-                },
-                "properties": {"Node name for S&R": "LayerMask: SegmentAnythingUltra V3"}
-            },
-            "10": {
-                "classType": "Image Seamless Texture",
-                "inputs": {
-                    "blending": 0.37,
-                    "images": ["17", 0],
-                    "tiled": "true",
-                    "tiles": 2
-                },
-                "properties": {"Node name for S&R": "Image Seamless Texture"}
-            },
-            "13": {
-                "classType": "Paste By Mask",
-                "inputs": {
-                    "image_base": ["2", 0],
-                    "image_to_paste": ["10", 0],
-                    "mask": ["8", 0],
-                    "resize_behavior": "resize"
-                },
-                "properties": {"Node name for S&R": "Paste By Mask"}
-            },
-            "17": {
-                "classType": "TensorArt_LoadImage",
-                "inputs": {
-                    "_height": 768,
-                    "_width": 512,
-                    "image": texture_resource_id,
-                    "upload": "image"
-                },
-                "properties": {"Node name for S&R": "TensorArt_LoadImage"}
-            },
-            "2": {
-                "classType": "TensorArt_LoadImage",
-                "inputs": {
-                    "_height": 1024,
-                    "_width": 768,
-                    "image": image_resource_id,
-                    "upload": "image"
-                },
-                "properties": {"Node name for S&R": "TensorArt_LoadImage"}
-            },
-            "3": {
-                "classType": "LayerMask: LoadSegmentAnythingModels",
-                "inputs": {
-                    "grounding_dino_model": "GroundingDINO_SwinB (938MB)",
-                    "sam_model": "sam_vit_h (2.56GB)"
-                },
-                "properties": {"Node name for S&R": "LayerMask: LoadSegmentAnythingModels"}
-            },
-            "4": {
-                "classType": "TensorArt_PromptText",
-                "inputs": {"Text": position.lower()},
-                "properties": {"Node name for S&R": "TensorArt_PromptText"}
-            },
-            "7": {
-                "classType": "PreviewImage",
-                "inputs": {"images": ["13", 0]},
-                "properties": {"Node name for S&R": "PreviewImage"}
-            },
-            "8": {
-                "classType": "MaskToImage",
-                "inputs": {"mask": ["1", 1]},
-                "properties": {"Node name for S&R": "MaskToImage"}
-            }
-        }
-
-        payload = {
-            "requestId": f"workflow_{int(time.time())}",
-            "params": workflow_params,
-            "runningNotifyUrl": ""
+            "1": { "classType": "LayerMask: SegmentAnythingUltra V3", "inputs": {"image": ["2", 0], "prompt": ["4", 0]}},
+            "2": {"classType": "TensorArt_LoadImage", "inputs": {"image": image_resource_id}},
+            "4": {"classType": "TensorArt_PromptText", "inputs": {"Text": position.lower()}},
+            "17": {"classType": "TensorArt_LoadImage", "inputs": {"image": texture_resource_id}},
+            # Thêm các node khác nếu cần từ mã gốc
         }
-        
-        output_path = run_workflow(payload, "full_workflow")
-        return output_path
-
+        payload = {"requestId": f"workflow_{int(time.time())}", "params": workflow_params}
+        return run_workflow(payload, "full_workflow")
     except Exception as e:
         print(f"Mask generation error: {str(e)}")
         return None
 
-
 def run_workflow(payload, step_name):
     headers = {
         'Content-Type': 'application/json',
@@ -389,5 +273,5 @@ with gr.Blocks() as demo:
 # Khởi chạy Gradio với API
 demo.launch(
     server_name="0.0.0.0",
-    server_port=7860
+    server_port=7860,
 )