Update app.py

app.py

@@ -5,9 +5,11 @@ from PIL import Image
 import requests
 from io import BytesIO
 
+# Replicate API client setup
+# Set environment variable: export RAPI_TOKEN="your_token_here"
 replicate_client = replicate.Client(api_token=os.getenv("RAPI_TOKEN"))
 
-# Comprehensive hairstyle options
+# Comprehensive hairstyle options - All 93 styles
 HAIRCUT_OPTIONS = [
     "Crew Cut", "Faux Hawk", "Slicked Back", "Side-Parted", "Center-Parted",
     "Blunt Bangs", "Side-Swept Bangs", "Shag", "Lob", "Angled Bob",
@@ -63,33 +65,18 @@ HAIR_COLOR_OPTIONS = {
     "Rose Gold": "#E0BFB8"
 }
 
-def create_color_button(color_name, color_code):
-    """Create HTML for a single color button"""
-    return f'''
-    <button class="color-btn" onclick="selectColor('{color_name}')" style="background-color: {color_code};" title="{color_name}">
-        <span class="color-label">{color_name}</span>
-    </button>
-    '''
-
-def create_color_palette():
-    """Create the complete color palette HTML"""
-    palette_html = '<div class="color-palette">'
-    for color_name, color_code in HAIR_COLOR_OPTIONS.items():
-        palette_html += create_color_button(color_name, color_code)
-    palette_html += '</div>'
-    return palette_html
-
-def change_haircut(input_image, haircut_style, hair_color_state):
+def change_haircut(input_image, haircut_style, hair_color):
     """
-    Process image through Replicate API to change hairstyle
+    Process image through RAPI to change hairstyle
     """
     try:
         # Check if image is provided
         if input_image is None:
-            return None, "Please upload an image first.", hair_color_state
+            return None, "Please upload an image first."
         
-        # Use the color from state
-        hair_color = hair_color_state if hair_color_state else "Black"
+        # Ensure we have a valid color
+        if not hair_color or hair_color not in HAIR_COLOR_OPTIONS:
+            hair_color = "Black"
         
         # Save PIL Image to temporary file
         temp_path = "temp_input.png"
@@ -119,36 +106,23 @@ def change_haircut(input_image, haircut_style, hair_color_state):
             if os.path.exists(temp_path):
                 os.remove(temp_path)
                 
-            return result_image, f"Successfully applied {haircut_style} style with {hair_color} color!", hair_color_state
+            return result_image, f"Successfully applied {haircut_style} style with {hair_color} color!"
         else:
-            return None, "Processing error occurred.", hair_color_state
+            return None, "Processing error occurred."
             
     except Exception as e:
         # Clean up temporary file on error
         if os.path.exists(temp_path):
             os.remove(temp_path)
-        return None, f"Error: {str(e)}", hair_color_state
+        return None, f"Error: {str(e)}"
+
+def update_color_display(color):
+    """Update the color display text"""
+    return f"**Selected Color:** {color}"
 
 # Create Gradio interface
 with gr.Blocks(title="AI Hairstyle Changer", theme=gr.themes.Soft()) as demo:
-    # Hidden state for selected color
-    selected_color = gr.State(value="Black")
-    
-    gr.Markdown(
-        """
-        # 🎨 Ginigen AI Hairstyle Changer
-        
-        Transform your hairstyle with AI! Upload a photo and choose from 90+ hairstyle options and 29 hair colors.
-        
-        **Instructions:**
-        1. Upload a front-facing photo with clear visibility of your face
-        2. Select your desired hairstyle from the options below
-        3. Click on a hair color from the palette
-        4. Click "Apply Hairstyle" to see the transformation
-        
-        ⚠️ **Note:** Requires Replicate API key. Set REPLICATE_API_TOKEN in your environment.
-        """
-    )
+
     
     with gr.Row():
         with gr.Column(scale=1):
@@ -162,7 +136,23 @@ with gr.Blocks(title="AI Hairstyle Changer", theme=gr.themes.Soft()) as demo:
             with gr.Group():
                 gr.Markdown("### Select Hair Color")
                 selected_color_display = gr.Markdown(value="**Selected Color:** Black")
-                color_palette = gr.HTML(create_color_palette())
+                
+                # Create color buttons in a grid
+                color_buttons = []
+                with gr.Column():
+                    # Create rows of color buttons
+                    for i in range(0, len(HAIR_COLOR_OPTIONS), 5):
+                        with gr.Row():
+                            for color_name, color_code in list(HAIR_COLOR_OPTIONS.items())[i:i+5]:
+                                btn = gr.Button(
+                                    value=color_name,
+                                    elem_classes="color-button",
+                                    elem_id=f"color-{color_name.replace(' ', '-')}"
+                                )
+                                color_buttons.append((btn, color_name, color_code))
+                
+                # Hidden textbox to store selected color
+                selected_color = gr.Textbox(value="Black", visible=False)
             
             submit_btn = gr.Button(
                 "🎨 Apply Hairstyle",
@@ -173,13 +163,13 @@ with gr.Blocks(title="AI Hairstyle Changer", theme=gr.themes.Soft()) as demo:
         with gr.Column(scale=2):
             # Hairstyle selection with scrollable radio buttons
             with gr.Group():
-                gr.Markdown("### Select Hairstyle")
+                gr.Markdown(f"### Select Hairstyle (Total: {len(HAIRCUT_OPTIONS)} styles)")
                 with gr.Column(elem_classes="hairstyle-container"):
                     haircut_radio = gr.Radio(
                         choices=HAIRCUT_OPTIONS,
                         value="Bob",
                         label="Hairstyle Options",
-                        info="Browse and select from 90+ hairstyle options",
+                        info="Browse and select from 93 hairstyle options",
                         elem_classes="hairstyle-radio"
                     )
         
@@ -197,36 +187,8 @@ with gr.Blocks(title="AI Hairstyle Changer", theme=gr.themes.Soft()) as demo:
                 value="Upload a photo and select a hairstyle to begin."
             )
     
-    # JavaScript for color selection
-    demo.load(None, None, None, js="""
-    function() {
-        window.selectColor = function(colorName) {
-            // Update all color buttons to remove selected class
-            document.querySelectorAll('.color-btn').forEach(btn => {
-                btn.classList.remove('selected');
-            });
-            
-            // Add selected class to clicked button
-            event.target.closest('.color-btn').classList.add('selected');
-            
-            // Update the display text
-            const displayElement = document.querySelector('#component-5 > div.prose.markdown-text');
-            if (displayElement) {
-                displayElement.innerHTML = '<strong>Selected Color:</strong> ' + colorName;
-            }
-            
-            // Store the selected color
-            window.selectedHairColor = colorName;
-        };
-        
-        // Set initial selection
-        window.selectedHairColor = 'Black';
-        document.querySelector('.color-btn').classList.add('selected');
-    }
-    """)
-    
-    # Custom CSS
-    demo.css = """
+    # Custom CSS with dynamic color styling
+    custom_css = """
     .hairstyle-container {
         max-height: 600px;
         overflow-y: auto;
@@ -249,77 +211,92 @@ with gr.Blocks(title="AI Hairstyle Changer", theme=gr.themes.Soft()) as demo:
         background-color: #e8e8e8;
     }
     
-    .color-palette {
-        display: grid;
-        grid-template-columns: repeat(auto-fill, minmax(90px, 1fr));
-        gap: 8px;
-        padding: 10px;
-        max-height: 300px;
-        overflow-y: auto;
-        border: 1px solid #e0e0e0;
-        border-radius: 8px;
-        background-color: #f9f9f9;
+    .color-button {
+        min-height: 50px !important;
+        margin: 2px !important;
+        border: 3px solid transparent !important;
+        border-radius: 8px !important;
+        font-weight: bold !important;
+        color: white !important;
+        text-shadow: 1px 1px 2px rgba(0,0,0,0.7) !important;
+        transition: all 0.3s ease !important;
+        position: relative !important;
     }
     
-    .color-btn {
-        position: relative;
-        width: 100%;
-        height: 50px;
-        border: 3px solid transparent;
-        border-radius: 8px;
-        cursor: pointer;
-        transition: all 0.3s ease;
-        overflow: hidden;
-        box-shadow: 0 2px 5px rgba(0,0,0,0.2);
-    }
-    
-    .color-btn:hover {
-        transform: translateY(-2px);
-        box-shadow: 0 4px 10px rgba(0,0,0,0.3);
+    .color-button:hover {
+        transform: translateY(-2px) !important;
+        box-shadow: 0 4px 10px rgba(0,0,0,0.3) !important;
     }
     
-    .color-btn.selected {
-        border-color: #2196F3;
-        transform: scale(1.05);
-        box-shadow: 0 0 0 3px rgba(33, 150, 243, 0.3);
+    .color-button.selected {
+        border-color: #2196F3 !important;
+        transform: scale(1.05) !important;
+        box-shadow: 0 0 0 3px rgba(33, 150, 243, 0.3) !important;
     }
+    """
     
-    .color-label {
-        position: absolute;
-        bottom: 0;
-        left: 0;
-        right: 0;
-        background-color: rgba(0,0,0,0.7);
-        color: white;
-        font-size: 10px;
-        padding: 2px;
-        text-align: center;
-        font-weight: bold;
-    }
+    # Add dynamic CSS for each color button
+    for color_name, color_code in HAIR_COLOR_OPTIONS.items():
+        custom_css += f"""
+    #color-{color_name.replace(' ', '-')} > button {{
+        background-color: {color_code} !important;
+    }}
     """
     
-    # Button click event with JavaScript integration
-    def process_with_js_color(input_image, haircut_style):
-        """Process using color from JavaScript"""
-        return change_haircut(input_image, haircut_style, selected_color.value)
+    demo.css = custom_css
+    
+    # Connect color button clicks
+    for btn, color_name, color_code in color_buttons:
+        btn.click(
+            fn=lambda c=color_name: (c, update_color_display(c)),
+            outputs=[selected_color, selected_color_display],
+            js=f"""(x) => {{
+                // Remove selected class from all buttons
+                document.querySelectorAll('.color-button').forEach(b => b.classList.remove('selected'));
+                // Add selected class to clicked button
+                document.querySelector('#color-{color_name.replace(' ', '-')} > button').classList.add('selected');
+                return '{color_name}';
+            }}"""
+        )
     
+    # Button click event
     submit_btn.click(
-        fn=process_with_js_color,
-        inputs=[input_image, haircut_radio],
-        outputs=[output_image, status_text, selected_color],
-        js="() => { return [null, null, window.selectedHairColor || 'Black']; }"
-    ).then(
         fn=change_haircut,
         inputs=[input_image, haircut_radio, selected_color],
-        outputs=[output_image, status_text, selected_color]
+        outputs=[output_image, status_text]
+    )
+    
+    # Set initial selection
+    demo.load(
+        js="""() => {
+            document.querySelector('#color-Black > button').classList.add('selected');
+        }"""
     )
     
+    # Additional information
+    gr.Markdown(
+        """
+        ### 📚 Hairstyle Categories
+        - **Short Styles**: Pixie Cut, Crew Cut, Undercut, Mohawk
+        - **Medium Styles**: Bob, Lob, Shag, Layered
+        - **Long Styles**: Straight, Wavy, Curly
+        - **Braided Styles**: French Braid, Dutch Braid, Fishtail Braid, Box Braids
+        - **Updo Styles**: Chignon, French Twist, Messy Bun, Top Knot
+        - **Special Styles**: Space Buns, Dreadlocks, Cornrows, Beehive
+        
+        ### 🎨 Hair Color Categories
+        - **Natural Colors**: Black, Browns, Blonde variations
+        - **Red Tones**: Red, Auburn, Copper, Mahogany, Burgundy
+        - **Cool Tones**: Silver, Ash Blonde, Ash Brown, Titanium
+        - **Fashion Colors**: Blue, Purple, Pink, Green, Rose Gold
+        """
+    )
 
 # Run the app
 if __name__ == "__main__":
     # Check API key
     if not os.getenv("RAPI_TOKEN"):
-        print("⚠️  Warning: REAPI_TOKEN environment variable not set!")
+        print("⚠️  Warning: RAPI_TOKEN environment variable not set!")
         print("Set it using: export RAPI_TOKEN='your_token'")
     
     demo.launch(