Update app.py

app.py

@@ -1,11 +1,8 @@
 import gradio as gr
-import numpy as np
 import json
 import requests
 import os
-from typing import Dict, List, Tuple
 from PIL import Image, ImageDraw
-import io
 
 # Fireworks AI configuration
 FIREWORKS_API_KEY = os.getenv("FIREWORKS_API_KEY", "YOUR_API_KEY_HERE")
@@ -19,7 +16,7 @@ BODY_PARTS = {
     "LEye": 15, "REar": 16, "LEar": 17
 }
 
-# Skeleton connections for drawing
+# Skeleton connections
 POSE_CONNECTIONS = [
     ("Neck", "RShoulder"), ("RShoulder", "RElbow"), ("RElbow", "RWrist"),
     ("Neck", "LShoulder"), ("LShoulder", "LElbow"), ("LElbow", "LWrist"),
@@ -82,7 +79,7 @@ POSE_TEMPLATES = {
     }
 }
 
-def draw_pose(keypoints: Dict, width: int = 512, height: int = 512) -> Image.Image:
+def draw_pose(keypoints, width=512, height=512):
     """Draw pose skeleton on image"""
     img = Image.new('RGB', (width, height), color='white')
     draw = ImageDraw.Draw(img)
@@ -101,12 +98,10 @@ def draw_pose(keypoints: Dict, width: int = 512, height: int = 512) -> Image.Ima
             x, y = point
             radius = 5
             draw.ellipse([x-radius, y-radius, x+radius, y+radius], fill='red', outline='darkred')
-            # Add label
-            draw.text((x+8, y-8), part[:3], fill='black')
     
     return img
 
-def generate_pose_from_llm(prompt: str) -> Dict:
+def generate_pose_from_llm(prompt):
     """Generate pose using LLM"""
     system_prompt = """You are an expert in generating human pose keypoints.
     Given a description, generate 18 keypoint coordinates for OpenPose.
@@ -116,9 +111,7 @@ def generate_pose_from_llm(prompt: str) -> Dict:
                RHip, RKnee, RAnkle, LHip, LKnee, LAnkle, REye, LEye, REar, LEar
     
     Return ONLY a JSON object with keypoint names and [x, y] coordinates.
-    Example: {"Nose": [256, 80], "Neck": [256, 120], ...}
-    
-    Ensure anatomically correct proportions and center the pose."""
+    Example: {"Nose": [256, 80], "Neck": [256, 120], ...}"""
     
     headers = {
         "Accept": "application/json",
@@ -142,7 +135,6 @@ def generate_pose_from_llm(prompt: str) -> Dict:
             data = response.json()
             content = data['choices'][0]['message']['content']
             
-            # Extract JSON from response
             import re
             json_match = re.search(r'\{.*\}', content, re.DOTALL)
             if json_match:
@@ -151,11 +143,13 @@ def generate_pose_from_llm(prompt: str) -> Dict:
     except Exception as e:
         print(f"LLM Error: {e}")
     
-    # Fallback to template
     return get_template_from_prompt(prompt)
 
-def get_template_from_prompt(prompt: str) -> Dict:
-    """Select appropriate template based on prompt"""
+def get_template_from_prompt(prompt):
+    """Select template based on prompt"""
+    if not prompt:
+        return POSE_TEMPLATES["Standing"]
+    
     prompt_lower = prompt.lower()
     
     if any(word in prompt_lower for word in ["sit", "chair", "seated"]):
@@ -171,56 +165,41 @@ def get_template_from_prompt(prompt: str) -> Dict:
     else:
         return POSE_TEMPLATES["Standing"]
 
-def refine_pose(current_keypoints: Dict, instruction: str) -> Dict:
-    """Refine existing pose based on instruction"""
+def refine_pose(current_keypoints, instruction):
+    """Refine existing pose"""
+    if not current_keypoints or not instruction:
+        return current_keypoints
+    
     keypoints = current_keypoints.copy()
     instruction_lower = instruction.lower()
     
-    # Simple rule-based refinement
+    # Simple refinements
     if "raise" in instruction_lower or "lift" in instruction_lower:
-        if "arm" in instruction_lower or "hand" in instruction_lower:
+        if "arm" in instruction_lower:
             if "left" in instruction_lower:
                 if "LWrist" in keypoints:
                     keypoints["LWrist"][1] -= 50
-                if "LElbow" in keypoints:
-                    keypoints["LElbow"][1] -= 30
             elif "right" in instruction_lower:
                 if "RWrist" in keypoints:
                     keypoints["RWrist"][1] -= 50
-                if "RElbow" in keypoints:
-                    keypoints["RElbow"][1] -= 30
-            else:  # Both arms
+            else:
                 for part in ["LWrist", "RWrist"]:
                     if part in keypoints:
                         keypoints[part][1] -= 50
-                for part in ["LElbow", "RElbow"]:
-                    if part in keypoints:
-                        keypoints[part][1] -= 30
     
     elif "lower" in instruction_lower:
-        if "arm" in instruction_lower or "hand" in instruction_lower:
+        if "arm" in instruction_lower:
             for part in ["LWrist", "RWrist"]:
                 if part in keypoints:
                     keypoints[part][1] += 50
     
-    elif "spread" in instruction_lower or "wide" in instruction_lower:
-        if "leg" in instruction_lower:
-            if "LAnkle" in keypoints:
-                keypoints["LAnkle"][0] -= 30
-            if "RAnkle" in keypoints:
-                keypoints["RAnkle"][0] += 30
-    
-    elif "bend" in instruction_lower:
-        if "knee" in instruction_lower:
-            for part in ["LKnee", "RKnee"]:
-                if part in keypoints:
-                    keypoints[part][1] += 20
-                    keypoints[part][0] += 10 if "L" in part else -10
-    
     return keypoints
 
-def keypoints_to_openpose_format(keypoints: Dict) -> str:
-    """Convert keypoints to OpenPose JSON format"""
+def keypoints_to_openpose_format(keypoints):
+    """Convert to OpenPose JSON format"""
+    if not keypoints:
+        return "{}"
+    
     candidate = []
     for i in range(18):
         part_name = None
@@ -240,231 +219,115 @@ def keypoints_to_openpose_format(keypoints: Dict) -> str:
     
     return json.dumps({"candidate": candidate, "subset": subset}, indent=2)
 
+# Main generation function
+def generate_pose(prompt, use_llm, template):
+    """Main function to generate pose"""
+    if template and template != "None":
+        keypoints = POSE_TEMPLATES[template]
+    elif use_llm and FIREWORKS_API_KEY != "YOUR_API_KEY_HERE" and prompt:
+        keypoints = generate_pose_from_llm(prompt)
+    elif prompt:
+        keypoints = get_template_from_prompt(prompt)
+    else:
+        keypoints = POSE_TEMPLATES["Standing"]
+    
+    pose_img = draw_pose(keypoints)
+    json_str = keypoints_to_openpose_format(keypoints)
+    
+    return pose_img, json_str, keypoints
+
+def refine_existing_pose(instruction, keypoints_json):
+    """Refine pose with instruction"""
+    if not keypoints_json:
+        return None, "{}", {}
+    
+    refined_keypoints = refine_pose(keypoints_json, instruction)
+    pose_img = draw_pose(refined_keypoints)
+    json_str = keypoints_to_openpose_format(refined_keypoints)
+    
+    return pose_img, json_str, refined_keypoints
+
+def check_api_status():
+    """Check if API key is configured"""
+    if FIREWORKS_API_KEY != "YOUR_API_KEY_HERE":
+        return "✅ API key configured - Advanced AI ready"
+    return "⚠️ API key not set - Template mode active"
+
 # Create Gradio interface
-def create_demo():
-    with gr.Blocks(title="AI Pose Generator", theme=gr.themes.Soft()) as demo:
-        current_keypoints = gr.State({})
-        
-        gr.Markdown("""
-        # 🎨 AI Line Art Pose Generator
-        ### Generate precise poses from text descriptions using AI
-        """)
-        
-        with gr.Tabs():
-            with gr.TabItem("🤖 Generate Pose"):
-                with gr.Row():
-                    with gr.Column(scale=1):
-                        # Input section
-                        use_llm = gr.Checkbox(
-                            label="🚀 Use Advanced AI Model (Fireworks API)", 
-                            value=False,
-                            info="Enable for more accurate pose generation (requires API key)"
-                        )
-                        
-                        api_status = gr.Markdown("⚠️ API key not set - Template mode active")
-                        
-                        prompt = gr.Textbox(
-                            label="Pose Description",
-                            placeholder="e.g., A person sitting cross-legged reading a book",
-                            lines=3
-                        )
-                        
-                        gr.Examples(
-                            examples=[
-                                "A person standing with arms raised in victory",
-                                "Someone sitting at a desk typing on a laptop",
-                                "A dancer in arabesque position",
-                                "A person doing yoga warrior pose",
-                                "Someone waving hello",
-                                "A person running with arms pumping"
-                            ],
-                            inputs=prompt
-                        )
-                        
-                        generate_btn = gr.Button("🎯 Generate Pose", variant="primary", size="lg")
-                        
-                        # Template selection
-                        with gr.Accordion("📚 Quick Templates", open=False):
-                            template_select = gr.Dropdown(
-                                choices=list(POSE_TEMPLATES.keys()),
-                                label="Select Template",
-                                value="Standing"
-                            )
-                            use_template_btn = gr.Button("Apply Template")
-                    
-                    with gr.Column(scale=1):
-                        # Output section
-                        pose_image = gr.Image(
-                            label="Generated Pose",
-                            type="pil",
-                            height=512
-                        )
-                        
-                        with gr.Accordion("📋 OpenPose JSON", open=False):
-                            json_output = gr.Code(
-                                label="JSON Data",
-                                language="json",
-                                lines=10
-                            )
-            
-            with gr.TabItem("✏️ Refine Pose"):
-                with gr.Row():
-                    with gr.Column():
-                        refinement_instruction = gr.Textbox(
-                            label="Refinement Instructions",
-                            placeholder="e.g., Raise the left arm higher",
-                            lines=2
-                        )
-                        
-                        refine_btn = gr.Button("✨ Refine Pose", variant="secondary")
-                        
-                        gr.Markdown("""
-                        **Quick commands:**
-                        - "Raise left/right arm"
-                        - "Lower arms"
-                        - "Spread legs wider"
-                        - "Bend knees"
-                        """)
-                        
-                        # Manual adjustment
-                        with gr.Accordion("🎛️ Manual Adjustment", open=False):
-                            selected_part = gr.Dropdown(
-                                choices=list(BODY_PARTS.keys()),
-                                label="Select Body Part",
-                                value="RWrist"
-                            )
-                            x_adjust = gr.Slider(-100, 100, 0, label="X Adjustment")
-                            y_adjust = gr.Slider(-100, 100, 0, label="Y Adjustment")
-                            apply_adjust_btn = gr.Button("Apply Adjustment")
-                    
-                    with gr.Column():
-                        refined_image = gr.Image(
-                            label="Refined Pose",
-                            type="pil",
-                            height=512
-                        )
-                        
-                        with gr.Accordion("📋 Updated JSON", open=False):
-                            refined_json = gr.Code(
-                                label="JSON Data",
-                                language="json",
-                                lines=10
-                            )
+with gr.Blocks(title="AI Pose Generator") as demo:
+    keypoints_state = gr.State({})
+    
+    gr.Markdown("""
+    # 🎨 AI Line Art Pose Generator
+    ### Generate precise poses from text descriptions
+    """)
+    
+    with gr.Row():
+        with gr.Column(scale=1):
+            # Input controls
+            api_status = gr.Markdown(check_api_status())
             
-            with gr.TabItem("ℹ️ Help"):
-                gr.Markdown("""
-                ## How to Use
-                
-                ### 1. Generate Pose
-                - Enter a natural language description of the pose
-                - Click "Generate Pose" to create the pose
-                - Or select a template for quick start
-                
-                ### 2. Refine Pose (Optional)
-                - Use natural language commands to adjust the pose
-                - Or manually adjust individual body parts
-                
-                ### 3. Export
-                - Copy the OpenPose JSON format for use in other applications
-                - Compatible with ControlNet and other pose-based tools
-                
-                ### API Setup (Optional)
-                For better results, set up Fireworks API:
-                ```bash
-                export FIREWORKS_API_KEY="your_api_key"
-                ```
-                
-                ### Features
-                - 🚀 No GPU required - runs on CPU
-                - 🎨 Clean line art style
-                - 📊 OpenPose compatible format
-                - 🔧 Easy refinement tools
-                - 💾 JSON export for integration
-                """)
-        
-        # Event handlers
-        def check_api_status():
-            if FIREWORKS_API_KEY != "YOUR_API_KEY_HERE":
-                return "✅ API key configured - Advanced AI ready"
-            return "⚠️ API key not set - Template mode active"
-        
-        def generate_pose(prompt_text, use_llm_flag):
-            if not prompt_text:
-                keypoints = POSE_TEMPLATES["Standing"]
-            elif use_llm_flag and FIREWORKS_API_KEY != "YOUR_API_KEY_HERE":
-                keypoints = generate_pose_from_llm(prompt_text)
-            else:
-                keypoints = get_template_from_prompt(prompt_text)
+            use_llm = gr.Checkbox(
+                label="Use Advanced AI (Fireworks API)",
+                value=False
+            )
             
-            pose_img = draw_pose(keypoints)
-            json_str = keypoints_to_openpose_format(keypoints)
+            prompt = gr.Textbox(
+                label="Describe the pose",
+                placeholder="e.g., A person sitting and reading a book",
+                lines=2
+            )
             
-            return pose_img, json_str, keypoints
-        
-        def use_template(template_name):
-            keypoints = POSE_TEMPLATES[template_name]
-            pose_img = draw_pose(keypoints)
-            json_str = keypoints_to_openpose_format(keypoints)
-            return pose_img, json_str, keypoints
-        
-        def refine_existing_pose(instruction, keypoints_state):
-            if not keypoints_state:
-                gr.Warning("Please generate a pose first")
-                return None, None, keypoints_state
+            template = gr.Dropdown(
+                choices=["None"] + list(POSE_TEMPLATES.keys()),
+                label="Or select a template",
+                value="None"
+            )
             
-            refined_keypoints = refine_pose(keypoints_state, instruction)
-            pose_img = draw_pose(refined_keypoints)
-            json_str = keypoints_to_openpose_format(refined_keypoints)
-            return pose_img, json_str, refined_keypoints
-        
-        def manual_adjust(part, x_adj, y_adj, keypoints_state):
-            if not keypoints_state:
-                gr.Warning("Please generate a pose first")
-                return None, None, keypoints_state
+            generate_btn = gr.Button("🎯 Generate Pose", variant="primary")
             
-            if part not in keypoints_state:
-                gr.Warning(f"Part {part} not found in current pose")
-                return None, None, keypoints_state
+            # Refinement
+            gr.Markdown("### Refine Pose")
+            refinement = gr.Textbox(
+                label="Refinement instruction",
+                placeholder="e.g., Raise the left arm",
+                lines=1
+            )
             
-            adjusted_keypoints = keypoints_state.copy()
-            adjusted_keypoints[part][0] += x_adj
-            adjusted_keypoints[part][1] += y_adj
+            refine_btn = gr.Button("✨ Refine", variant="secondary")
             
-            pose_img = draw_pose(adjusted_keypoints)
-            json_str = keypoints_to_openpose_format(adjusted_keypoints)
-            return pose_img, json_str, adjusted_keypoints
+            # Examples
+            gr.Examples(
+                examples=[
+                    "A person standing with arms raised",
+                    "Someone sitting at a desk",
+                    "A person doing yoga",
+                    "Someone waving hello",
+                    "A person running"
+                ],
+                inputs=prompt
+            )
         
-        # Connect events
-        demo.load(check_api_status, outputs=api_status)
-        
-        generate_btn.click(
-            generate_pose,
-            inputs=[prompt, use_llm],
-            outputs=[pose_image, json_output, current_keypoints]
-        )
-        
-        use_template_btn.click(
-            use_template,
-            inputs=[template_select],
-            outputs=[pose_image, json_output, current_keypoints]
-        )
-        
-        refine_btn.click(
-            refine_existing_pose,
-            inputs=[refinement_instruction, current_keypoints],
-            outputs=[refined_image, refined_json, current_keypoints]
-        )
-        
-        apply_adjust_btn.click(
-            manual_adjust,
-            inputs=[selected_part, x_adjust, y_adjust, current_keypoints],
-            outputs=[refined_image, refined_json, current_keypoints]
-        )
-        
-        return demo
-
-# Create and launch the app
-app = create_demo()
+        with gr.Column(scale=1):
+            # Output
+            pose_image = gr.Image(label="Generated Pose", type="pil")
+            
+            with gr.Accordion("OpenPose JSON", open=False):
+                json_output = gr.Code(language="json", lines=10)
+    
+    # Event handlers
+    generate_btn.click(
+        fn=generate_pose,
+        inputs=[prompt, use_llm, template],
+        outputs=[pose_image, json_output, keypoints_state]
+    )
+    
+    refine_btn.click(
+        fn=refine_existing_pose,
+        inputs=[refinement, keypoints_state],
+        outputs=[pose_image, json_output, keypoints_state]
+    )
 
+# Launch the app
 if __name__ == "__main__":
-    app.launch()
+    demo.launch()