app.py

import os
import base64
import requests
from flask import Flask, request, jsonify, send_from_directory
from flask_cors import CORS
from PIL import Image
import io
from dotenv import load_dotenv

# Load environment variables from .env file if present
load_dotenv()

# Get environment variables
PORT = int(os.environ.get('PORT', 7860))
HOST = os.environ.get('HOST', '0.0.0.0')
GROK_IMAGE_MODEL = os.environ.get('GROK_IMAGE_MODEL', 'grok-2-image-1212')
GROK_VISION_MODEL = os.environ.get('GROK_VISION_MODEL', 'grok-2-vision-1212')

app = Flask(__name__, static_folder='static')
CORS(app)

# Route to serve the frontend
@app.route('/', defaults={'path': ''})
@app.route('/<path:path>')
def serve(path):
    if path != "" and os.path.exists(os.path.join(app.static_folder, path)):
        return send_from_directory(app.static_folder, path)
    else:
        return send_from_directory(app.static_folder, 'index.html')

# Image generation API
@app.route('/api/generate-image', methods=['POST'])
def generate_image():
    data = request.json
    api_key = data.get('api_key')
    prompt = data.get('prompt')
    use_openai_format = data.get('use_openai_format', False)
    
    if not api_key or not prompt:
        return jsonify({"error": "Missing API key or prompt"}), 400
    
    try:
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {api_key}"
        }
        
        # OpenAI SDK-style format or original format - both use the same endpoint
        json_data = {
            "model": "grok-2-image-1212",  # Use the specific Grok-2 image model
            "prompt": prompt
            # Simplified request: removed n and size parameters
        }
        
        print(f"Sending request to x.ai API with data: {json_data}")
        print(f"Using authorization header: Bearer {api_key[:10]}...")
        
        # The correct endpoint for x.ai image generation
        response = requests.post(
            "https://api.x.ai/v1/images/generations",
            headers=headers,
            json=json_data
        )
        
        if not response.ok:
            # Print detailed error information
            print(f"X.AI API Error Status: {response.status_code}")
            print(f"X.AI API Error Response: {response.text}")
            try:
                error_data = response.json()
                print(f"X.AI API Error JSON: {error_data}")
                return jsonify({"error": error_data.get('error', {}).get('message', f"API error: {response.status_code}")}), response.status_code
            except:
                return jsonify({"error": f"API error: {response.status_code}, Response: {response.text}"}), response.status_code
        
        result = response.json()
        print(f"X.AI API Success Response: {result}")
        
        # Handle different possible response formats
        image_url = None
        caption = None
        
        if 'data' in result and len(result['data']) > 0:
            data_item = result['data'][0]
            if 'url' in data_item:
                image_url = data_item['url']
            if 'revised_prompt' in data_item:
                caption = data_item['revised_prompt']
        elif 'images' in result and len(result['images']) > 0:
            image_url = result['images'][0]
        elif 'url' in result:
            image_url = result['url']
        
        if not image_url:
            print(f"No image URL found in response: {result}")
            return jsonify({"error": "Unexpected API response format", "details": result}), 500
            
        response_data = {"image_url": image_url}
        if caption:
            response_data["caption"] = caption
            
        return jsonify(response_data)
        
    except requests.exceptions.RequestException as e:
        # If the API returns an error response
        print(f"Request exception: {str(e)}")
        if hasattr(e, 'response') and e.response is not None:
            try:
                error_data = e.response.json()
                print(f"Error response JSON: {error_data}")
                return jsonify({"error": error_data.get('error', {}).get('message', str(e))}), e.response.status_code
            except:
                print(f"Could not parse error response: {e.response.text}")
                return jsonify({"error": f"API error: {e.response.status_code}"}), e.response.status_code
        return jsonify({"error": str(e)}), 500

# Image analysis API
@app.route('/api/analyze-image', methods=['POST'])
def analyze_image():
    api_key = request.form.get('api_key')
    prompt = request.form.get('prompt', 'Describe this image in detail')
    
    if not api_key or 'image' not in request.files:
        return jsonify({"error": "Missing API key or image"}), 400
    
    try:
        # Process the uploaded image
        image_file = request.files['image']
        img = Image.open(image_file)
        
        # Convert to base64
        buffered = io.BytesIO()
        img.save(buffered, format=img.format if img.format else 'JPEG')
        img_str = base64.b64encode(buffered.getvalue()).decode('utf-8')
        
        # Determine image mime type
        format_mapping = {
            'JPEG': 'image/jpeg',
            'PNG': 'image/png',
            'GIF': 'image/gif',
            'WEBP': 'image/webp'
        }
        mime_type = format_mapping.get(img.format, 'image/jpeg')
        
        # Call the Grok Vision API
        response = requests.post(
            "https://api.x.ai/v1/chat/completions",
            headers={
                "Content-Type": "application/json",
                "Authorization": f"Bearer {api_key}"
            },
            json={
                "model": GROK_VISION_MODEL,
                "messages": [
                    {
                        "role": "user",
                        "content": [
                            {"type": "text", "text": prompt},
                            {
                                "type": "image_url",
                                "image_url": {
                                    "url": f"data:{mime_type};base64,{img_str}"
                                }
                            }
                        ]
                    }
                ],
                "max_tokens": 1000
            }
        )
        
        response.raise_for_status()
        result = response.json()
        
        # Handle different response formats
        if 'choices' in result and len(result['choices']) > 0 and 'message' in result['choices'][0]:
            analysis = result['choices'][0]['message']['content']
        elif 'response' in result:
            analysis = result['response']
        else:
            return jsonify({"error": "Unexpected API response format", "details": result}), 500
            
        return jsonify({"analysis": analysis})
        
    except requests.exceptions.RequestException as e:
        # If the API returns an error response
        if hasattr(e, 'response') and e.response is not None:
            try:
                error_data = e.response.json()
                return jsonify({"error": error_data.get('error', {}).get('message', str(e))}), e.response.status_code
            except:
                return jsonify({"error": f"API error: {e.response.status_code}"}), e.response.status_code
        return jsonify({"error": str(e)}), 500
    except Exception as e:
        return jsonify({"error": str(e)}), 500

if __name__ == '__main__':
    print(f"Starting Grok-2 API interface on {HOST}:{PORT}")
    print(f"Using image model: {GROK_IMAGE_MODEL}")
    print(f"Using vision model: {GROK_VISION_MODEL}")
    app.run(host=HOST, port=PORT)