Spaces:
Sleeping
Sleeping
| import os | |
| import base64 | |
| import json | |
| from io import BytesIO | |
| from typing import List, Optional, Union, Dict, Any | |
| from fastapi import FastAPI, File, UploadFile, HTTPException, Depends, Security, status, Form | |
| from fastapi.security import APIKeyHeader | |
| from pydantic import BaseModel, Field | |
| import requests | |
| import uvicorn | |
| from PIL import Image | |
| app = FastAPI() | |
| # Environment variables for authentication | |
| HF_TOKEN = os.getenv("HF_TOKEN") | |
| API_KEY = os.getenv("API_KEY") | |
| if not HF_TOKEN: | |
| raise ValueError("Hugging Face API token (HF_TOKEN) is missing!") | |
| if not API_KEY: | |
| raise ValueError("API key (API_KEY) is missing!") | |
| # API Key security scheme | |
| api_key_header = APIKeyHeader(name="X-API-Key", auto_error=True) | |
| # Authentication dependency | |
| async def get_api_key(api_key: str = Security(api_key_header)): | |
| if api_key != API_KEY: | |
| raise HTTPException( | |
| status_code=status.HTTP_401_UNAUTHORIZED, | |
| detail="Invalid API Key" | |
| ) | |
| return api_key | |
| # Models | |
| class Message(BaseModel): | |
| role: str | |
| content: Union[str, List[Dict[str, Any]]] | |
| def to_dict(self): | |
| return {"role": self.role, "content": self.content} | |
| class ChatItem(BaseModel): | |
| messages: List[Message] | |
| temperature: float = 0.7 | |
| max_tokens: int = 1024 | |
| top_p: float = 0.9 | |
| frequency_penalty: float = 0.0 | |
| presence_penalty: float = 0.0 | |
| class ImageUploadRequest(BaseModel): | |
| prompt: str | |
| system_prompt: Optional[str] = "" | |
| temperature: float = 0.7 | |
| max_tokens: int = 1024 | |
| top_p: float = 0.9 | |
| def encode_image_to_base64(image_file): | |
| """Convert an image file to base64 string""" | |
| img = Image.open(image_file) | |
| buffer = BytesIO() | |
| img.save(buffer, format=img.format or "JPEG") | |
| img_str = base64.b64encode(buffer.getvalue()).decode("utf-8") | |
| return img_str | |
| def generate_with_phi_vision(messages_data): | |
| """Generate response using Phi-3.5-Vision model""" | |
| headers = { | |
| "Authorization": f"Bearer {HF_TOKEN}", | |
| "Content-Type": "application/json" | |
| } | |
| api_url = "https://api-inference.huggingface.co/models/microsoft/Phi-3.5-vision-instruct" | |
| # Convert pydantic models to dictionaries if necessary | |
| message_dicts = [] | |
| for msg in messages_data: | |
| if hasattr(msg, 'to_dict'): | |
| message_dicts.append(msg.to_dict()) | |
| else: | |
| message_dicts.append(msg) | |
| payload = {"inputs": message_dicts} | |
| try: | |
| response = requests.post(api_url, json=payload, headers=headers) | |
| if response.status_code != 200: | |
| # Handle the "model is loading" scenario | |
| if "is currently loading" in response.text: | |
| error_data = response.json() | |
| raise HTTPException( | |
| status_code=503, | |
| detail=f"Model is currently loading. Estimated time: {error_data.get('estimated_time', 'unknown')} seconds" | |
| ) | |
| else: | |
| raise HTTPException( | |
| status_code=response.status_code, | |
| detail=f"Hugging Face API error: {response.text}" | |
| ) | |
| return response.json() | |
| except requests.exceptions.RequestException as e: | |
| raise HTTPException( | |
| status_code=status.HTTP_503_SERVICE_UNAVAILABLE, | |
| detail=f"Error connecting to Hugging Face API: {str(e)}" | |
| ) | |
| async def generate_text(item: ChatItem, api_key: str = Depends(get_api_key)): | |
| """Generate text response using Phi-3.5-Vision model""" | |
| try: | |
| # Convert Message objects to dictionaries | |
| messages_dict = [msg.to_dict() for msg in item.messages] | |
| response = generate_with_phi_vision(messages_dict) | |
| return {"response": response} | |
| except Exception as e: | |
| if isinstance(e, HTTPException): | |
| raise e | |
| raise HTTPException(status_code=500, detail=str(e)) | |
| async def generate_with_image( | |
| prompt: str = Form(...), | |
| system_prompt: str = Form(""), | |
| image: UploadFile = File(...), | |
| temperature: float = Form(0.7), | |
| max_tokens: int = Form(1024), | |
| top_p: float = Form(0.9), | |
| api_key: str = Depends(get_api_key) | |
| ): | |
| """Generate response from Phi-3.5-Vision based on an image and text prompt""" | |
| try: | |
| # Convert image to base64 | |
| img_base64 = encode_image_to_base64(image.file) | |
| # Determine content type | |
| content_type = image.content_type or "image/jpeg" | |
| # Construct the messages in the format expected by the model | |
| messages = [ | |
| { | |
| "role": "system", | |
| "content": system_prompt if system_prompt else "You are a helpful vision assistant." | |
| }, | |
| { | |
| "role": "user", | |
| "content": [ | |
| {"type": "text", "text": prompt}, | |
| {"type": "image_url", "image_url": {"url": f"data:{content_type};base64,{img_base64}"}} | |
| ] | |
| } | |
| ] | |
| # Call the model with retry logic | |
| response = generate_with_phi_vision(messages) | |
| return {"response": response} | |
| except Exception as e: | |
| if isinstance(e, HTTPException): | |
| raise e | |
| raise HTTPException(status_code=500, detail=str(e)) | |
| async def test_api(api_key: str = Depends(get_api_key)): | |
| """Test if the API is working""" | |
| return {"message": "API is working!"} | |
| if __name__ == "__main__": | |
| uvicorn.run(app, host="0.0.0.0", port=7860) |