Update app.py

app.py

@@ -15,21 +15,27 @@ api_key = os.getenv("HF_API_KEY")
 image_api_url = os.getenv("IMAGE_API_URL")
 headers = {"Authorization": f"Bearer {api_key}"}
 
-# Danh sách màu sắc và mã màu hex tương ứng
+# Danh sách màu sắc
 color_map = {
-    "Đỏ": "#FF0000",
-    "Cam": "#FF7F00",
-    "Vàng": "#FFFF00",
-    "Lục": "#00FF00",
-    "Lam": "#0000FF",
-    "Chàm": "#4B0082",
-    "Tím": "#800080",
-    "Trắng": "#FFFFFF",
-    "Đen": "#000000",
-    "Xám": "#808080",
-    "Cầu vồng": "#FF69B4",  # Ví dụ màu cầu vồng, có thể thay bằng màu sắc khác nếu cần
+    "Đỏ": "red",
+    "Cam": "orange",
+    "Vàng": "yellow",
+    "Lục": "green",
+    "Lam": "blue",
+    "Chàm": "indigo",
+    "Tím": "violet",
+    "Trắng": "white",
+    "Đen": "black",
+    "Xám": "gray",
+    "Cầu vồng": "rainbow",
 }
 
+# Danh sách kiểu tóc
+hair_styles = [
+    "Tóc tự nhiên", "Tóc thẳng mượt", "Tóc buộc nửa đầu", "Tóc búi cao", "Tóc bob", 
+    "Tóc xoăn sóng", "Tóc tết", "Tóc pixie"
+]
+
 # Hàm gọi API Hugging Face để tạo hình ảnh
 def query(payload):
     response = requests.post(image_api_url, headers=headers, json=payload)
@@ -42,17 +48,12 @@ def translate_to_english(text):
     return translate(text, "en", "vi")  # Dịch từ tiếng Việt sang tiếng Anh
 
 # Hàm xử lý tạo hình ảnh với các lựa chọn
-def generate_image(prompt, view, style, hair_color, eye_color, lip_color):
+def generate_image(prompt, view, style, hair_color, eye_color, lip_color, hair_style):
     # Dịch prompt từ tiếng Việt sang tiếng Anh
     prompt_en = translate_to_english(prompt)
     
-    # Lấy mã màu hex tương ứng từ tên màu đã chọn
-    hair_color_hex = color_map.get(hair_color, "#FF0000")
-    eye_color_hex = color_map.get(eye_color, "#FF0000")
-    lip_color_hex = color_map.get(lip_color, "#FF0000")
-    
     # Thêm các yếu tố lựa chọn vào prompt
-    prompt_with_choices = f"{prompt_en}, View: {view}, Style: {style}, Hair color: {hair_color_hex}, Eye color: {eye_color_hex}, Lip color: {lip_color_hex}"
+    prompt_with_choices = f"{prompt_en}, View: {view}, Style: {style}, Hair color: {hair_color}, Eye color: {eye_color}, Lip color: {lip_color}, Hair style: {hair_style}"
 
     # Thêm yếu tố ngẫu nhiên vào prompt
     random_factor = random.randint(1, 1000)
@@ -84,6 +85,10 @@ iface = gr.Interface(
             ["Đỏ", "Cam", "Vàng", "Lục", "Lam", "Chàm", "Tím", "Trắng", "Đen", "Xám", "Cầu vồng"],
             label="Màu môi"
         ),
+        gr.Dropdown(
+            hair_styles, 
+            label="Kiểu tóc"
+        ),
     ],
     outputs="image",
     title="Image Generator - Trần Như Tuấn",