import gradio as gr import requests import json import os import hashlib import time from PIL import Image from pathlib import Path from io import BytesIO from groq import Groq from dotenv import load_dotenv # Load environment variables load_dotenv() api_key_token = os.getenv('api_key_token', '') groq_api_key = os.getenv('groq_api_key', '') if not api_key_token or not groq_api_key: raise ValueError("Please configure your API keys in .env") # Khởi tạo client Groq (sửa lỗi proxies) original_init = Groq.__init__ def patched_init(self, *args, **kwargs): kwargs.pop("proxies", None) # Loại bỏ tham số proxies nếu có original_init(self, *args, **kwargs) Groq.__init__ = patched_init client = Groq(api_key=groq_api_key) # Cấu hình cơ bản 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 PRODUCT_GROUPS = { "Standard": { "C1012 Glacier White": "817687427545199895", "C1026 Polar": "819910519797326073", "C3269 Ash Grey": "821839484099264081", "C3168 Silver Wave": "821849044696643212", "C1005 Milky White": "821948258441171133", }, "Deluxe": { "C2103 Onyx Carrara": "827090618489513527", "C2104 Massa": "822075428127644644", "C3105 Casla Cloudy": "828912225788997963", "C3146 Casla Nova": "828013009961087650", "C2240 Marquin": "828085015087780649", }, # Thêm các nhóm khác nếu cần } PRODUCT_IMAGE_MAP = { "C1012 Glacier White": "product_images/C1012.jpg", "C1026 Polar": "product_images/C1026.jpg", "C3269 Ash Grey": "product_images/C3269.jpg", "C3168 Silver Wave": "product_images/C3168.jpg", "C1005 Milky White": "product_images/C1005.jpg", "C2103 Onyx Carrara": "product_images/C2103.jpg", "C2104 Massa": "product_images/C2104.jpg", "C3105 Casla Cloudy": "product_images/C3105.jpg", "C3146 Casla Nova": "product_images/C3146.jpg", "C2240 Marquin": "product_images/C2240.jpg", "C2262 Concrete (Honed)": "product_images/C2262.jpg", "C3311 Calacatta Sky": "product_images/C3311.jpg", "C3346 Massimo": "product_images/C3346.jpg", "C4143 Mario": "product_images/C4143.jpg", "C4145 Marina": "product_images/C4145.jpg", "C4202 Calacatta Gold": "product_images/C4202.jpg", "C1205 Casla Everest": "product_images/C1205.jpg", "C4211 Calacatta Supreme": "product_images/C4211.jpg", "C4204 Calacatta Classic": "product_images/C4204.jpg", "C1102 Super White": "product_images/C1102.jpg", "C4246 Casla Mystery": "product_images/C4246.jpg", "C4345 Oro": "product_images/C4345.jpg", "C4346 Luxe": "product_images/C4346.jpg", "C4342 Casla Eternal": "product_images/C4342.jpg", "C4221 Athena": "product_images/C4221.jpg", "C4255 Calacatta Extra": "product_images/C4255.jpg", } # Hàm xử lý từ mã gốc def rewrite_prompt_with_groq(vietnamese_prompt, product_codes): prompt = f"{vietnamese_prompt}, featuring {' and '.join(product_codes)} quartz marble" return prompt def upload_image_to_tensorart(image_path): try: url = f"{url_pre}/resource/image" payload = json.dumps({"expireSec": "7200"}) headers = { 'Content-Type': 'application/json', 'Accept': 'application/json', 'Authorization': f'Bearer {api_key_token}' } if not os.path.exists(image_path): return None response = requests.post(url, headers=headers, data=payload, timeout=30) response.raise_for_status() resource_response = response.json() put_url = resource_response.get('putUrl') headers_put = resource_response.get('headers', {'Content-Type': 'image/jpeg'}) if not put_url: return None with open(image_path, 'rb') as img_file: upload_response = requests.put(put_url, data=img_file, headers=headers_put) if upload_response.status_code not in [200, 203]: raise Exception(f"PUT failed with status {upload_response.status_code}") resource_id = resource_response.get('resourceId') if not resource_id: return None time.sleep(10) # Đợi đồng bộ tài nguyên return resource_id except Exception as e: print(f"Upload error: {str(e)}") return None def txt2img(prompt, width, height, product_codes): model_id = "779398605850080514" vae_id = "ae.sft" txt2img_data = { "request_id": hashlib.md5(str(int(time.time())).encode()).hexdigest(), "stages": [ {"type": "INPUT_INITIALIZE", "inputInitialize": {"seed": -1, "count": 1}}, { "type": "DIFFUSION", "diffusion": { "width": width, "height": height, "prompts": [{"text": prompt}], "negativePrompts": [{"text": " "}], "sdModel": model_id, "sdVae": vae_id, "sampler": "Euler a", "steps": 30, "cfgScale": 8, "clipSkip": 1, "etaNoiseSeedDelta": 31337, } } ] } headers = { 'Content-Type': 'application/json', 'Accept': 'application/json', 'Authorization': f'Bearer {api_key_token}' } response = requests.post(f"{url_pre}/jobs", json=txt2img_data, headers=headers) if response.status_code != 200: return f"Error: {response.status_code} - {response.text}" response_data = response.json() job_id = response_data['job']['id'] start_time = time.time() timeout = 300 while True: time.sleep(10) elapsed_time = time.time() - start_time if elapsed_time > timeout: return f"Error: Job timed out after {timeout} seconds." response = requests.get(f"{url_pre}/jobs/{job_id}", headers=headers) if response.status_code != 200: return f"Error: {response.status_code} - {response.text}" get_job_response_data = response.json() job_status = get_job_response_data['job']['status'] if job_status == 'SUCCESS': image_url = get_job_response_data['job']['successInfo']['images'][0]['url'] response_image = requests.get(image_url) img = Image.open(BytesIO(response_image.content)) save_path = Path(SAVE_DIR) / f"{hashlib.md5(prompt.encode()).hexdigest()}.png" img.save(save_path) return save_path elif job_status == 'FAILED': return "Error: Job failed." def generate_mask(image_resource_id, position, selected_product_code): try: if not image_resource_id: 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) if not texture_filepath or not os.path.exists(texture_filepath): raise Exception(f"Texture file not found for {short_code}") texture_resource_id = upload_image_to_tensorart(texture_filepath) if not texture_resource_id: raise Exception(f"Failed to upload texture for {short_code}") time.sleep(10) workflow_params = { "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 } 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', 'Accept': 'application/json', 'Authorization': f'Bearer {api_key_token}' } response = requests.post(f"{url_pre}/jobs/workflow", json=payload, headers=headers, timeout=300) if response.status_code != 200: return None job_id = response.json()['job']['id'] max_attempts = 36 for _ in range(max_attempts): response = requests.get(f"{url_pre}/jobs/{job_id}", headers=headers, timeout=30) result = response.json() if result['job']['status'] == 'SUCCESS': image_url = result['job']['successInfo']['images'][0]['url'] image_response = requests.get(image_url) image = Image.open(BytesIO(image_response.content)) if image.mode == 'RGBA': image = image.convert('RGB') output_path = Path(SAVE_DIR) / f"{step_name}_{int(time.time())}.jpg" image.save(output_path) return str(output_path) time.sleep(5) return None # API endpoints cho frontend def api_text2img(prompt, size, custom_size, product_codes): try: width, height = map(int, (custom_size if size == "Custom size" else size).split("x")) if not product_codes: return {"error": "At least one product code required"} short_codes = [code.split()[0] for code in product_codes] rewritten_prompt = rewrite_prompt_with_groq(prompt, short_codes) result = txt2img(rewritten_prompt, width, height, short_codes) if isinstance(result, str): return {"error": result} return {"image_url": f"/file={result}"} except Exception as e: return {"error": str(e)} def api_img2img(file, position, size, custom_size, product_codes): try: width, height = map(int, (custom_size if size == "Custom size" else size).split("x")) if not product_codes: return {"error": "At least one product code required"} image_path = Path(SAVE_DIR) / f"input_{int(time.time())}.jpg" file.save(image_path) image_resource_id = upload_image_to_tensorart(str(image_path)) if not image_resource_id: return {"error": "Failed to upload image"} output_path = generate_mask(image_resource_id, position, product_codes[0]) if not output_path: return {"error": "Failed to generate image"} return {"image_url": f"/file={output_path}"} except Exception as e: return {"error": str(e)} # Giao diện Gradio (tùy chọn) with gr.Blocks() as demo: gr.Markdown("## Backend Gradio cho CaslaQuartz") gr.Interface(fn=lambda x: "API only", inputs="text", outputs="text") # Khởi chạy Gradio với API demo.launch( server_name="0.0.0.0", server_port=7860, )