Update app.py

app.py

@@ -8,15 +8,19 @@ from dotenv import load_dotenv
 # Load environment variables from .env file
 load_dotenv()
 
-# API Configuration - Hugging Face Spaces stores secrets in environment variables
+# API Configuration
 API_KEY = os.getenv('STABLECOG_API_KEY', '')
 if not API_KEY:
-    # For local testing fallback
     API_KEY = "StableCogKey"
 
 API_HOST = 'https://api.stablecog.com'
-API_ENDPOINT = '/v1/credits'
-API_URL = f'{API_HOST}{API_ENDPOINT}'
+
+# API Endpoints
+CREDITS_ENDPOINT = '/v1/credits'
+MODELS_ENDPOINT = '/v1/image/generation/models'
+
+CREDITS_URL = f'{API_HOST}{CREDITS_ENDPOINT}'
+MODELS_URL = f'{API_HOST}{MODELS_ENDPOINT}'
 
 headers = {
     'Authorization': f'Bearer {API_KEY}',
@@ -26,7 +30,7 @@ headers = {
 def check_credits():
     """Check StableCog credits and return formatted results"""
     try:
-        response = requests.get(API_URL, headers=headers, timeout=10)
+        response = requests.get(CREDITS_URL, headers=headers, timeout=10)
         
         if response.status_code == 200:
             res_json = response.json()
@@ -62,10 +66,9 @@ def check_credits():
             
             # Use total_remaining_credits from API or calculate it
             if total_remaining_credits == 0 and credits_list:
-                # Calculate total remaining from all credit entries
                 total_remaining_credits = sum(credit.get('remaining_amount', 0) for credit in credits_list)
             
-            # Calculate total credits (max of initial total or used + remaining)
+            # Calculate total credits
             total_credits = max(total_initial_credits, total_remaining_credits + total_used_credits)
             
             # Calculate percentage used
@@ -118,6 +121,85 @@ def check_credits():
             "message": str(e)
         }
 
+def get_available_models():
+    """Get available StableCog models"""
+    try:
+        response = requests.get(MODELS_URL, headers=headers, timeout=10)
+        
+        if response.status_code == 200:
+            res_json = response.json()
+            models_list = res_json.get('models', [])
+            
+            # Organize models by type
+            organized_models = []
+            for model in models_list:
+                # Extract model information
+                model_info = {
+                    'id': model.get('id', ''),
+                    'name': model.get('name', 'Unknown'),
+                    'description': model.get('description', ''),
+                    'type': model.get('type', 'unknown'),
+                    'is_public': model.get('is_public', False),
+                    'is_default': model.get('is_default', False),
+                    'is_community': model.get('is_community', False),
+                    'created_at': model.get('created_at', ''),
+                    'updated_at': model.get('updated_at', '')
+                }
+                organized_models.append(model_info)
+            
+            # Sort models: default/public first, then by name
+            organized_models.sort(key=lambda x: (
+                not x['is_default'],
+                not x['is_public'],
+                x['name'].lower()
+            ))
+            
+            # Count by type
+            model_count = len(organized_models)
+            public_count = sum(1 for m in organized_models if m['is_public'])
+            default_count = sum(1 for m in organized_models if m['is_default'])
+            community_count = sum(1 for m in organized_models if m['is_community'])
+            
+            result = {
+                "success": True,
+                "models": organized_models,
+                "total_models": model_count,
+                "public_models": public_count,
+                "default_models": default_count,
+                "community_models": community_count,
+                "raw_data": json.dumps(res_json, indent=2),
+                "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S UTC")
+            }
+            
+            return result
+            
+        else:
+            return {
+                "success": False,
+                "error": f"API Error: {response.status_code}",
+                "message": response.text if response.text else "No response text",
+                "status_code": response.status_code
+            }
+            
+    except requests.exceptions.Timeout:
+        return {
+            "success": False,
+            "error": "Timeout Error",
+            "message": "The request timed out. Please try again."
+        }
+    except requests.exceptions.ConnectionError:
+        return {
+            "success": False,
+            "error": "Connection Error",
+            "message": "Could not connect to the API. Check your internet connection."
+        }
+    except Exception as e:
+        return {
+            "success": False,
+            "error": f"Unexpected Error: {type(e).__name__}",
+            "message": str(e)
+        }
+
 def update_display():
     """Update the UI with credit information"""
     result = check_credits()
@@ -244,6 +326,106 @@ def update_display():
         
         return html_content, f"Error: {result.get('error', 'Unknown error')}\n\nDetails: {result.get('message', '')}", 0, 0
 
+def display_models():
+    """Display available models in a formatted way"""
+    result = get_available_models()
+    
+    if result["success"]:
+        models_html = ""
+        model_counter = 0
+        
+        for model in result["models"]:
+            model_counter += 1
+            
+            # Determine model type badge
+            model_type = model['type'].upper() if model['type'] else 'UNKNOWN'
+            if model['is_community']:
+                model_type_badge = f"<span style='background: #FF6B6B; padding: 2px 8px; border-radius: 12px; font-size: 11px; margin-left: 5px;'>COMMUNITY</span>"
+            elif model['is_default']:
+                model_type_badge = f"<span style='background: #4CAF50; padding: 2px 8px; border-radius: 12px; font-size: 11px; margin-left: 5px;'>DEFAULT</span>"
+            elif model['is_public']:
+                model_type_badge = f"<span style='background: #2196F3; padding: 2px 8px; border-radius: 12px; font-size: 11px; margin-left: 5px;'>PUBLIC</span>"
+            else:
+                model_type_badge = f"<span style='background: #9E9E9E; padding: 2px 8px; border-radius: 12px; font-size: 11px; margin-left: 5px;'>{model_type}</span>"
+            
+            # Model ID (shortened)
+            model_id_short = model['id'][:8] + "..." if len(model['id']) > 8 else model['id']
+            
+            models_html += f"""
+            <div style='background: linear-gradient(135deg, rgba(255,255,255,0.1) 0%, rgba(255,255,255,0.05) 100%); border-radius: 10px; padding: 15px; margin-bottom: 12px; border-left: 4px solid #667eea;'>
+                <div style='display: flex; justify-content: space-between; align-items: center; margin-bottom: 8px;'>
+                    <h3 style='margin: 0; font-size: 16px;'>
+                        #{model_counter}. {model['name']}
+                        {model_type_badge}
+                    </h3>
+                    <span style='font-size: 12px; opacity: 0.7;'>ID: {model_id_short}</span>
+                </div>
+                
+                <p style='margin: 0 0 10px 0; font-size: 14px; opacity: 0.9;'>{model['description'] or 'No description available'}</p>
+                
+                <div style='display: flex; gap: 10px; font-size: 12px; opacity: 0.7;'>
+                    <span>📅 Created: {model['created_at'][:10] if model['created_at'] else 'Unknown'}</span>
+                    <span>🔄 Updated: {model['updated_at'][:10] if model['updated_at'] else 'Unknown'}</span>
+                </div>
+            </div>
+            """
+        
+        # Create stats section
+        stats_html = f"""
+        <div style='display: grid; grid-template-columns: repeat(2, 1fr); gap: 10px; margin-bottom: 20px;'>
+            <div style='background: rgba(102, 126, 234, 0.2); padding: 15px; border-radius: 8px; text-align: center;'>
+                <div style='font-size: 28px; font-weight: bold;'>{result['total_models']}</div>
+                <div style='font-size: 12px; opacity: 0.8;'>Total Models</div>
+            </div>
+            <div style='background: rgba(76, 175, 80, 0.2); padding: 15px; border-radius: 8px; text-align: center;'>
+                <div style='font-size: 28px; font-weight: bold;'>{result['public_models']}</div>
+                <div style='font-size: 12px; opacity: 0.8;'>Public Models</div>
+            </div>
+            <div style='background: rgba(33, 150, 243, 0.2); padding: 15px; border-radius: 8px; text-align: center;'>
+                <div style='font-size: 28px; font-weight: bold;'>{result['default_models']}</div>
+                <div style='font-size: 12px; opacity: 0.8;'>Default Models</div>
+            </div>
+            <div style='background: rgba(255, 107, 107, 0.2); padding: 15px; border-radius: 8px; text-align: center;'>
+                <div style='font-size: 28px; font-weight: bold;'>{result['community_models']}</div>
+                <div style='font-size: 12px; opacity: 0.8;'>Community Models</div>
+            </div>
+        </div>
+        """
+        
+        html_content = f"""
+        <div style='font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, sans-serif; padding: 25px; background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); border-radius: 15px; color: white; box-shadow: 0 10px 30px rgba(0,0,0,0.2);'>
+            <h2 style='text-align: center; margin-bottom: 25px; font-size: 28px; font-weight: 600;'>🤖 Available StableCog Models</h2>
+            
+            <div style='background: rgba(255,255,255,0.1); backdrop-filter: blur(10px); padding: 25px; border-radius: 12px; margin-bottom: 20px; border: 1px solid rgba(255,255,255,0.2);'>
+                {stats_html}
+                <div style='max-height: 500px; overflow-y: auto; padding-right: 10px;'>
+                    {models_html}
+                </div>
+            </div>
+            
+            <div style='text-align: center; font-size: 14px; opacity: 0.7; margin-top: 10px;'>
+                ⏰ Last updated: {result['timestamp']} | 🔄 Refresh to see latest models
+            </div>
+        </div>
+        """
+        
+        return html_content, result['raw_data']
+    
+    else:
+        # Error display
+        html_content = f"""
+        <div style='font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, sans-serif; padding: 25px; background: linear-gradient(135deg, #f093fb 0%, #f5576c 100%); border-radius: 15px; color: white; text-align: center; box-shadow: 0 10px 30px rgba(0,0,0,0.2);'>
+            <h2 style='margin-bottom: 20px; font-size: 26px;'>⚠️ Failed to Load Models</h2>
+            
+            <div style='background: rgba(255,255,255,0.15); backdrop-filter: blur(10px); padding: 20px; border-radius: 10px; margin-bottom: 20px; border: 1px solid rgba(255,255,255,0.2);'>
+                <p style='margin: 0 0 10px 0; font-size: 18px; font-weight: bold;'>{result.get('error', 'Unknown error')}</p>
+                <p style='margin: 0; font-size: 14px; opacity: 0.9;'>{result.get('message', '')}</p>
+            </div>
+        </div>
+        """
+        
+        return html_content, f"Error: {result.get('error', 'Unknown error')}\n\nDetails: {result.get('message', '')}"
+
 def get_recommendation(remaining_credits, percentage_used):
     """Provide recommendations based on credit status"""
     if remaining_credits == 0:
@@ -259,84 +441,93 @@ def get_recommendation(remaining_credits, percentage_used):
     else:
         return "🚀 **Plenty of credits!** Ready for extensive image generation work and experimentation."
 
-# Create Gradio interface with custom theme
-theme = gr.themes.Soft(
-    primary_hue="purple",
-    secondary_hue="indigo",
-    font=[gr.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui", "sans-serif"]
-).set(
-    button_primary_background_fill="linear-gradient(135deg, #667eea 0%, #764ba2 100%)",
-    button_primary_background_fill_hover="linear-gradient(135deg, #764ba2 0%, #667eea 100%)",
-    button_primary_text_color="white",
-    button_primary_border_color="rgba(255,255,255,0.2)",
-)
-
-with gr.Blocks(theme=theme, title="StableCog Credit Monitor", css="footer {display: none !important;}") as demo:
+# Create Gradio interface with Gradio 6 theme
+# Note: Gradio 6 has a simplified theme system
+with gr.Blocks(
+    theme=gr.themes.Soft(
+        primary_hue="purple",
+        secondary_hue="indigo",
+        font=[gr.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui", "sans-serif"]
+    ),
+    title="StableCog Dashboard",
+    css="""
+    footer {display: none !important;}
+    .gradio-container {max-width: 1400px !important;}
+    .tab-nav {background: rgba(255,255,255,0.1) !important; border-radius: 10px !important; padding: 5px !important;}
+    .stat-box input {font-weight: bold !important; font-size: 18px !important;}
+    .recommendation-box textarea {font-size: 16px !important; line-height: 1.5 !important;}
+    """
+) as demo:
     gr.Markdown("""
-    # 🎯 StableCog Credit Dashboard
+    # 🎯 StableCog Dashboard
     
-    *Monitor your API credits and plan your image generation projects efficiently.*
+    *Monitor your credits and explore available AI models for image generation.*
     """)
     
-    # Status row
-    with gr.Row():
-        with gr.Column(scale=2):
-            html_output = gr.HTML(label="Credit Status")
-        with gr.Column(scale=1):
-            raw_output = gr.Code(
-                label="📋 Raw API Response", 
-                language="json", 
-                interactive=False,
-                lines=15
-            )
-    
-    # Stats row
-    with gr.Row():
-        with gr.Column():
-            credits_display = gr.Number(
-                label="Remaining Credits", 
-                interactive=False,
-                elem_classes="stat-box"
-            )
-        with gr.Column():
-            usage_display = gr.Number(
-                label="Usage Percentage", 
-                interactive=False,
-                elem_classes="stat-box"
-            )
-    
-    # Recommendation row
-    with gr.Row():
-        recommendation_box = gr.Textbox(
-            label="🎯 AI Recommendation", 
-            interactive=False,
-            lines=3,
-            elem_classes="recommendation-box"
-        )
-    
-    # Control row
-    with gr.Row():
-        check_btn = gr.Button(
-            "🔄 Check Credits Now", 
-            variant="primary", 
-            size="lg",
-            elem_classes="refresh-btn"
-        )
-    
-    # Auto-check on load
-    demo.load(fn=update_display, inputs=None, outputs=[html_output, raw_output, credits_display, usage_display])
-    
-    # Connect button and update recommendation
-    def update_all():
-        html, raw, credits, usage = update_display()
-        recommendation = get_recommendation(credits, usage)
-        return html, raw, credits, usage, recommendation
-    
-    check_btn.click(
-        fn=update_all, 
-        inputs=None, 
-        outputs=[html_output, raw_output, credits_display, usage_display, recommendation_box]
-    )
+    with gr.Tabs() as tabs:
+        with gr.Tab("💰 Credits Dashboard", id="credits"):
+            # Status row
+            with gr.Row():
+                with gr.Column(scale=2):
+                    credits_html_output = gr.HTML(label="Credit Status")
+                with gr.Column(scale=1):
+                    credits_raw_output = gr.Code(
+                        label="📋 Raw API Response", 
+                        language="json", 
+                        interactive=False,
+                        lines=15
+                    )
+            
+            # Stats row
+            with gr.Row():
+                with gr.Column():
+                    credits_display = gr.Number(
+                        label="Remaining Credits", 
+                        interactive=False,
+                        elem_classes="stat-box"
+                    )
+                with gr.Column():
+                    usage_display = gr.Number(
+                        label="Usage Percentage", 
+                        interactive=False,
+                        elem_classes="stat-box"
+                    )
+            
+            # Recommendation row
+            with gr.Row():
+                recommendation_box = gr.Textbox(
+                    label="🎯 AI Recommendation", 
+                    interactive=False,
+                    lines=3,
+                    elem_classes="recommendation-box"
+                )
+            
+            # Control row
+            with gr.Row():
+                check_credits_btn = gr.Button(
+                    "🔄 Check Credits", 
+                    variant="primary",
+                    scale=0
+                )
+        
+        with gr.Tab("🤖 Available Models", id="models"):
+            with gr.Row():
+                with gr.Column(scale=2):
+                    models_html_output = gr.HTML(label="Available Models")
+                with gr.Column(scale=1):
+                    models_raw_output = gr.Code(
+                        label="📋 Raw API Response", 
+                        language="json", 
+                        interactive=False,
+                        lines=15
+                    )
+            
+            with gr.Row():
+                check_models_btn = gr.Button(
+                    "🔄 Refresh Models", 
+                    variant="primary",
+                    scale=0
+                )
     
     # Instructions and info
     with gr.Accordion("📚 How to Use & Setup", open=False):
@@ -349,33 +540,64 @@ with gr.Blocks(theme=theme, title="StableCog Credit Monitor", css="footer {displ
              - Key: `STABLECOG_API_KEY`
              - Value: `your_actual_api_key_here`
         
-        2. **Understanding Your Credits:**
-           - **Total Credits**: Sum of all initial credits received
-           - **Remaining Credits**: Currently available credits (sum of all credit types)
-           - **Used Credits**: Credits spent on image generation
-           - **Credit Types**: Different credit sources (Free, Refund, etc.)
-        
-        3. **Credit Breakdown:**
-           - **Free Credits**: Base credits provided to users
-           - **Refund Credits**: Credits returned for failed generations
-           - Each credit type may have different amounts and expiration
+        ### Features:
+        - **💰 Credits Dashboard**: Monitor your credit usage and remaining balance
+        - **🤖 Available Models**: Browse all StableCog image generation models
+        - **Smart Recommendations**: Get AI-powered suggestions based on your credits
+        - **Credit Breakdown**: See detailed breakdown by credit type
         
-        4. **Recommendations are based on:**
-           - Remaining credit count
-           - Usage percentage
-           - Common workload patterns
+        ### Understanding Credits:
+        - **Free Credits**: Base credits provided to users
+        - **Refund Credits**: Credits returned for failed generations
+        - Each credit type may have different amounts and expiration
         """)
     
     gr.Markdown("""
     ---
     *Built with ❤️ for StableCog users | [Report Issues](https://github.com/stability-ai/stablecog/issues)*
     """)
+    
+    # Initialize displays on load
+    demo.load(
+        fn=update_display, 
+        inputs=None, 
+        outputs=[credits_html_output, credits_raw_output, credits_display, usage_display]
+    )
+    
+    demo.load(
+        fn=display_models,
+        inputs=None,
+        outputs=[models_html_output, models_raw_output]
+    )
+    
+    # Update recommendation after credits load
+    def update_recommendation_after_load():
+        return get_recommendation(credits_display.value, usage_display.value)
+    
+    # Connect buttons
+    def update_all_credits():
+        html, raw, credits, usage = update_display()
+        recommendation = get_recommendation(credits, usage)
+        return html, raw, credits, usage, recommendation
+    
+    check_credits_btn.click(
+        fn=update_all_credits, 
+        inputs=None, 
+        outputs=[credits_html_output, credits_raw_output, credits_display, usage_display, recommendation_box]
+    )
+    
+    check_models_btn.click(
+        fn=display_models,
+        inputs=None,
+        outputs=[models_html_output, models_raw_output]
+    )
 
 if __name__ == "__main__":
-    # For Hugging Face Spaces
+    # For Hugging Face Spaces - Gradio 6 syntax
     demo.launch(
         server_name="0.0.0.0",
         server_port=7860,
         share=False,
-        show_error=True
+        show_error=True,
+        debug=False
     )