File size: 13,278 Bytes
c034a74
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
2bb03a8
 
ff704b5
 
 
d0ae17f
c034a74
6f1334b
c034a74
 
 
 
2bb03a8
 
 
 
 
 
ff704b5
c034a74
2bb03a8
 
3a240c4
c034a74
3a240c4
 
 
 
 
 
 
 
 
c034a74
2bb03a8
 
 
 
 
 
 
 
 
 
 
 
 
 
 
c034a74
2bb03a8
a318fb7
2bb03a8
 
 
 
 
a1a0caf
2bb03a8
 
3a240c4
2bb03a8
a318fb7
a1a0caf
2bb03a8
3a240c4
2bb03a8
 
3a240c4
 
2bb03a8
3a240c4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
2bb03a8
 
 
 
3a240c4
 
 
 
2bb03a8
 
 
3a240c4
 
 
 
2bb03a8
 
 
 
 
 
 
 
3a240c4
 
2bb03a8
a87cf29
3a240c4
 
 
 
a87cf29
 
3a240c4
 
 
 
 
 
 
 
 
 
 
 
 
2bb03a8
 
 
 
 
 
a318fb7
2bb03a8
 
908288f
e3f5ff0
2bb03a8
 
 
 
 
 
 
ff704b5
2bb03a8
 
 
 
 
c034a74
2bb03a8
 
e3f5ff0
2bb03a8
 
 
 
 
 
c034a74
2bb03a8
 
 
 
 
 
 
 
 
 
fce7c66
2bb03a8
 
 
 
fce7c66
3a240c4
 
 
 
 
2bb03a8
3a240c4
2bb03a8
 
 
3a240c4
2bb03a8
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
3a240c4
2bb03a8
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
d0ae17f
2bb03a8
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
# 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)

from fastapi import FastAPI, Response, HTTPException
from fastapi.responses import FileResponse, JSONResponse
from kokoro import KPipeline
import soundfile as sf
import os
import numpy as np
import torch
from huggingface_hub import InferenceClient
from pydantic import BaseModel
import base64
from io import BytesIO
from PIL import Image
import logging
from typing import Optional

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

class TextImageRequest(BaseModel):
    text: Optional[str] = None
    image_base64: Optional[str] = None
    voice: str = "af_heart"  # Default voice that we know exists
    speed: float = 1.0
    
    # List of known available voices - update this based on what's actually available
    AVAILABLE_VOICES = ["af_heart"]  # Add more voices as they become available
    
    # Validate that the voice exists
    def validate_voice(self):
        if self.voice not in self.AVAILABLE_VOICES:
            return "af_heart"  # Default to a voice we know exists
        return self.voice

class AudioResponse(BaseModel):
    status: str
    message: str

class ErrorResponse(BaseModel):
    error: str
    detail: Optional[str] = None

# Initialize FastAPI app
app = FastAPI(
    title="Text-to-Speech API with Vision Support",
    description="API for generating speech from text with optional image analysis",
    version="1.0.0"
)

def llm_chat_response(text, image_base64=None):
    """Function to get responses from LLM with text and optionally image input."""
    try:
        HF_TOKEN = os.getenv("HF_TOKEN")
        logger.info("Checking HF_TOKEN...")
        if not HF_TOKEN:
            logger.error("HF_TOKEN not found in environment variables")
            raise HTTPException(status_code=500, detail="HF_TOKEN not configured")
        
        logger.info("Initializing InferenceClient...")
        client = InferenceClient(
            provider="together",  # Using the provider shown in the sample
            api_key=HF_TOKEN
        )
        
        try:
            # IMPORTANT: Following exactly the format from the sample code
            if image_base64:
                logger.info("Processing request with image")
                prompt = text if text else "Describe this image in one sentence."
                
                messages = [
                    {
                        "role": "user",
                        "content": [
                            {
                                "type": "text",
                                "text": prompt
                            },
                            {
                                "type": "image_url",
                                "image_url": {
                                    "url": f"data:image/jpeg;base64,{image_base64}"
                                }
                            }
                        ]
                    }
                ]
            else:
                logger.info("Processing text-only request")
                messages = [
                    {
                        "role": "user",
                        "content": text + " Describe in one line only."
                    }
                ]
            
            logger.info("Sending request to model...")
            # Log the exact message structure we're sending
            logger.info(f"Message structure: {messages}")
            
            # Use the exact model name and parameters from the sample
            completion = client.chat.completions.create(
                model="meta-llama/Llama-3.2-11B-Vision-Instruct",
                messages=messages,
                max_tokens=500
            )
            
            logger.info(f"Received response from model")
            
            # Very simple response handling exactly like the sample code
            logger.info(f"Model response received: {completion}")
            
            try:
                # Extract response using the exact approach from the sample code
                response = completion.choices[0].message.content
                logger.info(f"Extracted response content: {response}")
                return response
            except Exception as e:
                logger.error(f"Error extracting message content: {str(e)}")
                logger.error(f"Attempting alternative extraction method...")
                
                # Fallback method if the above fails
                try:
                    if hasattr(completion.choices[0], "message"):
                        if hasattr(completion.choices[0].message, "content"):
                            return completion.choices[0].message.content
                    
                    # Last resort - try accessing as dictionary
                    return completion.choices[0]["message"]["content"]
                except Exception as e2:
                    logger.error(f"All extraction methods failed: {str(e2)}")
                    return "I couldn't process that input. Please try again with a different query."
            
        except Exception as e:
            logger.error(f"Error during model inference: {str(e)}")
            # Fallback response in case of error
            return "I couldn't process that input. Please try again with a different image or text query."
            
    except Exception as e:
        logger.error(f"Error in llm_chat_response: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

# Initialize pipeline once at startup
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)}")
    # We'll let the app start anyway, but log the error

@app.post("/generate", response_model=None, responses={
    200: {"content": {"application/octet-stream": {}}},
    400: {"model": ErrorResponse},
    500: {"model": ErrorResponse}
})
async def generate_audio(request: TextImageRequest):
    """
    Generate audio from text and optionally analyze an image.
    
    - If text is provided, uses that as input
    - If image is provided, analyzes the image
    - Converts the LLM response to speech using the specified voice and speed
    """
    try:
        logger.info(f"Received audio generation request")
        
        # If no text is provided but image is provided, use default prompt
        user_text = request.text if request.text is not None else ""
        if not user_text and request.image_base64:
            user_text = "Describe what you see in the image"
        elif not user_text and not request.image_base64:
            logger.error("Neither text nor image provided in request")
            return JSONResponse(
                status_code=400, 
                content={"error": "Request must include either text or image_base64"}
            )
        
        # Generate response using text and image if provided
        logger.info("Getting LLM response...")
        text_reply = llm_chat_response(user_text, request.image_base64)
        logger.info(f"LLM response: {text_reply}")
        
        # Validate voice parameter
        validated_voice = request.validate_voice()
        if validated_voice != request.voice:
            logger.warning(f"Requested voice '{request.voice}' not available, using '{validated_voice}' instead")
        
        # Generate audio
        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+'
            )
            
            # Process only the first segment for demo
            for i, (gs, ps, audio) in enumerate(generator):
                logger.info(f"Audio generated successfully: segment {i}")
                
                # 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"
                    }
                )
                
            logger.error("No audio segments generated")
            return JSONResponse(
                status_code=400,
                content={"error": "No audio generated", "detail": "The pipeline did not produce any audio"}
            )
            
        except Exception as e:
            logger.error(f"Error generating audio: {str(e)}")
            return JSONResponse(
                status_code=500,
                content={"error": "Audio generation failed", "detail": str(e)}
            )
    
    except Exception as e:
        logger.error(f"Unexpected error in generate_audio endpoint: {str(e)}")
        return JSONResponse(
            status_code=500,
            content={"error": "Internal server error", "detail": str(e)}
        )

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

@app.exception_handler(404)
async def not_found_handler(request, exc):
    return JSONResponse(
        status_code=404,
        content={"error": "Endpoint not found. Please use POST /generate for queries."}
    )

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