src/models/models.py

"""
Created By: ishwor subedi
Date: 2024-08-13
"""
from PIL import Image
from pydantic import BaseModel, Field
from typing import List, Dict


class LanguageTranslationRequest(BaseModel):
    text: str = Field(..., description="The text to translate")
    translator_backend_code: str = Field(..., description="The code for the translation backend")
    target: str = Field(..., description="The target language")


class ImageGenerationRequest(BaseModel):
    prompt: str = Field(..., description="The prompt for image generation")
    negative_prompt: str = Field(..., description="The negative prompt for image generation")
    style: str = Field(..., description="The style for image generation")
    use_negative_prompt: bool = Field(..., description="Whether to use the negative prompt")
    num_inference_steps: int = Field(..., description="The number of inference steps")
    num_images_per_prompt: int = Field(..., description="The number of images per prompt")
    seed: int = Field(..., description="The seed for image generation")
    width: int = Field(..., description="The width of the image")
    height: int = Field(..., description="The height of the image")
    guidance_scale: float = Field(..., description="The guidance scale for image generation")
    randomize_seed: bool = Field(..., description="Whether to randomize the seed")


class ImageCaptionRequest(BaseModel):
    prompt: str = Field(..., description="The prompt for image captioning")
    temperature: float = Field(..., description="The temperature for image captioning")
    length_penalty: float = Field(..., description="The length penalty for image captioning")
    repetition_penalty: float = Field(..., description="The repetition penalty for image captioning")
    max_length: int = Field(..., description="The maximum length for image captioning")
    min_length: int = Field(..., description="The minimum length for image captioning")
    top_p: float = Field(..., description="The top-p for image captioning")


class TextToSpeechRequest(BaseModel):
    text: str = Field(..., description="The text to convert to speech")
    lang: str = Field(..., description="The language of the text")
    tld: str = Field(..., description="The TLD of the language")


class SpeechToTextRequest(BaseModel):
    audio: str = Field(..., description="The audio to convert to text")
    lang: str = Field(..., description="The language of the audio")