app.py

# from fastapi import FastAPI, Response
# from fastapi.responses import FileResponse
# from kokoro import KPipeline
# import soundfile as sf
# import os
# import numpy as np
# import torch 
# from huggingface_hub import InferenceClient

# def llm_chat_response(text):
#     HF_TOKEN = os.getenv("HF_TOKEN")
#     client = InferenceClient(api_key=HF_TOKEN)
#     messages = [
# 	{
# 		"role": "user",
# 		"content": [
# 			{
# 				"type": "text",
# 				"text": text + str('describe in one line only')
# 			} #,
# 			# {
# 			# 	"type": "image_url",
# 			# 	"image_url": {
# 			# 		"url": "https://cdn.britannica.com/61/93061-050-99147DCE/Statue-of-Liberty-Island-New-York-Bay.jpg"
# 			# 	}
# 			# }
#             ]
# 	}
#     ]

#     response_from_llama = client.chat.completions.create(
#     model="meta-llama/Llama-3.2-11B-Vision-Instruct", 
# 	messages=messages, 
# 	max_tokens=500)

#     return response_from_llama.choices[0].message['content']

# app = FastAPI()

# # Initialize pipeline once at startup
# pipeline = KPipeline(lang_code='a')

# @app.post("/generate")
# async def generate_audio(text: str, voice: str = "af_heart", speed: float = 1.0):
    
#     text_reply = llm_chat_response(text)
    
#     # Generate audio
#     generator = pipeline(
#         text_reply, 
#         voice=voice,
#         speed=speed,
#         split_pattern=r'\n+'
#     )
    
#     # # Save first segment only for demo
#     # for i, (gs, ps, audio) in enumerate(generator):
#     #     sf.write(f"output_{i}.wav", audio, 24000)
#     #     return FileResponse(
#     #         f"output_{i}.wav",
#     #         media_type="audio/wav",
#     #         filename="output.wav"
#     #     )
    
#     # return Response("No audio generated", status_code=400)


#     # Process only the first segment for demo
#     for i, (gs, ps, audio) in enumerate(generator):

#         # Convert PyTorch tensor to NumPy array
#         audio_numpy = audio.cpu().numpy()
#         # Convert to 16-bit PCM
        
#         # Ensure the audio is in the range [-1, 1]
#         audio_numpy = np.clip(audio_numpy, -1, 1)
#         # Convert to 16-bit signed integers
#         pcm_data = (audio_numpy * 32767).astype(np.int16)
        
#         # Convert to bytes (automatically uses row-major order)
#         raw_audio = pcm_data.tobytes()
        
#         # Return PCM data with minimal necessary headers
#         return Response(
#             content=raw_audio,
#             media_type="application/octet-stream",
#             headers={
#                 "Content-Disposition": f'attachment; filename="output.pcm"',
#                 "X-Sample-Rate": "24000",
#                 "X-Bits-Per-Sample": "16",
#                 "X-Endianness": "little"
#             }
#         )
    
#     return Response("No audio generated", status_code=400)

import os
import uuid
import base64
import logging
from fastapi import FastAPI, HTTPException, Response, Request
from fastapi.responses import JSONResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
from typing import Optional, ClassVar, List
from huggingface_hub import InferenceClient
import numpy as np
import torch
from kokoro import KPipeline  # Your audio generation pipeline

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Create FastAPI app
app = FastAPI(
    title="Text-to-Speech API with Vision Support",
    description="This API uses meta-llama/Llama-3.2-11B-Vision-Instruct which requires an image input.",
    version="1.0.0"
)

# Mount a static directory for serving saved images
STATIC_DIR = "static_images"
if not os.path.exists(STATIC_DIR):
    os.makedirs(STATIC_DIR)
app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")

# Pydantic model for request
class TextImageRequest(BaseModel):
    text: Optional[str] = None
    image_base64: Optional[str] = None
    voice: str = "af_heart"  # Default voice
    speed: float = 1.0

    # Use ClassVar so that Pydantic doesn't treat this as a model field.
    AVAILABLE_VOICES: ClassVar[List[str]] = ["af_heart"]

    def validate_voice(self):
        if self.voice not in self.AVAILABLE_VOICES:
            return "af_heart"
        return self.voice

# Pydantic model for error responses
class ErrorResponse(BaseModel):
    error: str
    detail: Optional[str] = None

def llm_chat_response(prompt: str, image_base64: str) -> str:
    HF_TOKEN = os.getenv("HF_TOKEN")
    logger.info("Checking HF_TOKEN...")
    if not HF_TOKEN:
        logger.error("HF_TOKEN not configured")
        raise HTTPException(status_code=500, detail="HF_TOKEN not configured")
    
    logger.info("Initializing InferenceClient...")
    client = InferenceClient(
        provider="hf-inference",
        api_key=HF_TOKEN
    )
    
    # Save the base64-encoded image locally
    filename = f"{uuid.uuid4()}.jpg"
    image_path = os.path.join(STATIC_DIR, filename)
    try:
        image_data = base64.b64decode(image_base64)
    except Exception as e:
        logger.error(f"Error decoding image: {str(e)}")
        raise HTTPException(status_code=400, detail="Invalid base64 image data")
    
    with open(image_path, "wb") as f:
        f.write(image_data)
    
    # Construct the public URL for the saved image.
    # Set BASE_URL to your public URL if needed.
    base_url = os.getenv("BASE_URL", "http://localhost:8000")
    image_url = f"{base_url}/static/{filename}"
    
    # Build the message payload exactly as in the reference:
    messages = [
        {
            "role": "user",
            "content": [
                {
                    "type": "text",
                    "text": prompt
                },
                {
                    "type": "image_url",
                    "image_url": {
                        "url": image_url
                    }
                }
            ]
        }
    ]
    logger.info(f"Message structure: {messages}")
    
    try:
        completion = client.chat.completions.create(
            model="meta-llama/Llama-3.2-11B-Vision-Instruct",
            messages=messages,
            max_tokens=500,
        )
        response = completion.choices[0].message.content
        logger.info(f"Extracted response: {response}")
        return response
    except Exception as e:
        logger.error(f"Error during model inference: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

# Initialize the audio generation pipeline (KPipeline)
try:
    logger.info("Initializing KPipeline...")
    pipeline = KPipeline(lang_code='a')
    logger.info("KPipeline initialized successfully")
except Exception as e:
    logger.error(f"Failed to initialize KPipeline: {str(e)}")
    # The API will run but audio generation will fail if the pipeline is not ready.

@app.post("/generate", responses={
    200: {"content": {"application/octet-stream": {}}},
    400: {"model": ErrorResponse},
    500: {"model": ErrorResponse}
})
async def generate_audio(request: TextImageRequest):
    """
    Generate audio from a multimodal (text+image) input.
    This model requires an image input.
    """
    logger.info("Received generation request")
    
    # The model requires an image; if missing, return an error.
    if not request.image_base64:
        raise HTTPException(status_code=400, detail="This model requires an image input.")
    
    prompt = request.text if request.text else "Describe this image in one sentence."
    logger.info("Calling the LLM model")
    text_reply = llm_chat_response(prompt, request.image_base64)
    logger.info(f"LLM response: {text_reply}")
    
    validated_voice = request.validate_voice()
    if validated_voice != request.voice:
        logger.warning(f"Voice '{request.voice}' not available; using '{validated_voice}' instead")
    
    # Convert the text reply to audio using the KPipeline.
    logger.info(f"Generating audio using voice={validated_voice}, speed={request.speed}")
    try:
        generator = pipeline(
            text_reply,
            voice=validated_voice,
            speed=request.speed,
            split_pattern=r'\n+'
        )
        for _, _, audio in generator:
            audio_numpy = audio.cpu().numpy()
            audio_numpy = np.clip(audio_numpy, -1, 1)
            pcm_data = (audio_numpy * 32767).astype(np.int16)
            raw_audio = pcm_data.tobytes()
            return Response(
                content=raw_audio,
                media_type="application/octet-stream",
                headers={
                    "Content-Disposition": 'attachment; filename="output.pcm"',
                    "X-Sample-Rate": "24000",
                    "X-Bits-Per-Sample": "16",
                    "X-Endianness": "little"
                }
            )
        raise HTTPException(status_code=400, detail="No audio segments generated.")
    except Exception as e:
        logger.error(f"Error generating audio: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/")
async def root():
    return {"message": "Welcome to the Text-to-Speech API with Vision Support. Use POST /generate with text and image_base64."}

@app.exception_handler(404)
async def not_found_handler(request: Request, exc):
    return JSONResponse(status_code=404, content={"error": "Endpoint not found."})

@app.exception_handler(405)
async def method_not_allowed_handler(request: Request, exc):
    return JSONResponse(status_code=405, content={"error": "Method not allowed."})