app.py

import asyncio
import os
import torch
import redis
import io
import numpy as np
import base64
from PIL import Image
from dotenv import load_dotenv
from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments
from transformers import pipeline
from fastapi import FastAPI, HTTPException
from fastapi.responses import HTMLResponse
from typing import List, Dict, Any
import logging

# Configuración de logging
logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger(__name__)

# Cargar variables de entorno
load_dotenv()
HUGGINGFACE_TOKEN = os.getenv("HUGGINGFACE_TOKEN")
REDIS_HOST = os.getenv("REDIS_HOST")
REDIS_PORT = os.getenv("REDIS_PORT")
REDIS_PASSWORD = os.getenv("REDIS_PASSWORD")

# Configuración de Redis
redis_client = redis.StrictRedis(host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD, decode_responses=True)

# Inicializar la aplicación FastAPI
app = FastAPI()

# Diccionario para almacenar los modelos y sus propiedades en memoria
model_dict: Dict[str, Any] = {}
model_properties: Dict[str, Dict[str, Any]] = {}
model_lock = asyncio.Lock()

# Lista para almacenar el historial de mensajes y datos de entrenamiento en memoria
message_history: List[str] = []
TRAINING_DATA: List[Dict[str, torch.Tensor]] = []  # Datos de entrenamiento para texto
MUSIC_TRAINING_DATA: List[Dict[str, torch.Tensor]] = []  # Datos de entrenamiento para música
IMAGE_TRAINING_DATA: List[Dict[str, torch.Tensor]] = []  # Datos de entrenamiento para imágenes

# Parámetros adicionales para la generación de respuestas
TEMPERATURE = 0.7
TOP_PROBABILITY = 0.9
TOP_K = 50
FREQUENCY_PENALTY = 0.7
MAX_TOKENS = 1024  # Límite máximo de tokens por respuesta
UNIQUE_RESPONSES = set()  # Conjunto para almacenar respuestas únicas

# Inicializar el pipeline de generación de música y de imágenes
musicgen_pipeline = pipeline("text-to-audio", model="facebook/musicgen-small")
image_pipeline = pipeline("text-to-image", model="black-forest-labs/FLUX.1-schnell")

# Función para almacenar en Redis
def store_in_redis(key: str, value: Any):
    if isinstance(value, bytes):
        redis_client.set(key, value)
    else:
        redis_client.set(key, str(value))

# Función para recuperar de Redis
def retrieve_from_redis(key: str):
    value = redis_client.get(key)
    if value is None:
        return None
    try:
        return value.encode('latin1')  # Decodificar si es bytes
    except AttributeError:
        return value  # Si es texto

# Cargar modelos sincrónicamente al iniciar la aplicación
async def load_models():
    global model_dict, model_properties
    if not model_dict:  # Solo cargar si el diccionario está vacío
        for model_name in ["gpt2-medium", "gpt2-large", "gpt2", "google/gemma-2-9b", "meta-llama/Meta-Llama-3.1-8B-Instruct"]:
            model_key = f"model:{model_name}"
            tokenizer_key = f"tokenizer:{model_name}"
            
            model_data = retrieve_from_redis(model_key)
            tokenizer_data = retrieve_from_redis(tokenizer_key)
            
            if model_data and tokenizer_data:
                model = torch.load(io.BytesIO(model_data))
                tokenizer = torch.load(io.BytesIO(tokenizer_data))
                model_dict[model_name] = (model, tokenizer)
                logger.info(f"Loaded {model_name} from Redis")
            else:
                model = AutoModelForCausalLM.from_pretrained(model_name)
                tokenizer = AutoTokenizer.from_pretrained(model_name)
                model_dict[model_name] = (model, tokenizer)
                
                # Guardar modelos y tokenizers en Redis
                store_in_redis(model_key, torch.save(model, io.BytesIO()))
                store_in_redis(tokenizer_key, torch.save(tokenizer, io.BytesIO()))
                
                model_properties[model_name] = {
                    'pad_token': tokenizer.pad_token,
                    'pad_token_id': tokenizer.pad_token_id,
                    'eos_token': tokenizer.eos_token,
                    'eos_token_id': tokenizer.eos_token_id,
                    'bos_token': tokenizer.bos_token,
                    'bos_token_id': tokenizer.bos_token_id,
                    'unk_token': tokenizer.unk_token,
                    'unk_token_id': tokenizer.unk_token_id,
                    'padding_side': tokenizer.padding_side,
                    'special_tokens_map': tokenizer.special_tokens_map,
                    'model': model,
                    'tokenizer': tokenizer
                }
                logger.info(f"Successfully loaded {model_name} model and tokenizer")

# Cargar modelos una vez al iniciar la aplicación
asyncio.run(load_models())

# Funciones para las APIs adicionales de música e imágenes
def generate_music(prompt: str) -> bytes:
    # Generación de música utilizando el pipeline en memoria
    audio = musicgen_pipeline(prompt)['audio']
    return audio

def generate_image(prompt: str) -> bytes:
    # Generación de imagen utilizando el pipeline en memoria
    outputs = image_pipeline(prompt)["sample"][0]
    buffered = io.BytesIO()
    outputs.save(buffered, format="PNG")
    return buffered.getvalue()

# Ruta principal para la interfaz web
@app.get('/')
async def main():
    html_code = """
    <!DOCTYPE html>
    <html lang="en">
    <head>
        <meta charset="UTF-8">
        <meta http-equiv="X-UA-Compatible" content="IE=edge">
        <meta name="viewport" content="width=device-width, initial-scale=1.0">
        <title>ChatGPT Chatbot</title>
        <style>
            body, html {
                height: 100%;
                margin: 0;
                padding: 0;
                font-family: Arial, sans-serif;
            }
            .container {
                height: 100%;
                display: flex;
                flex-direction: column;
                justify-content: center;
                align-items: center;
            }
            .chat-container {
                border-radius: 10px;
                overflow: hidden;
                box-shadow: 0 0 10px rgba(0, 0, 0, 0.1);
                width: 100%;
                height: 100%;
                animation: fadeIn 1s ease;
            }
            .chat-box {
                height: calc(100% - 60px);
                overflow-y: auto;
                padding: 10px;
            }
            .chat-input {
                width: calc(100% - 100px);
                padding: 10px;
                border: none;
                border-top: 1px solid #ccc;
                font-size: 16px;
                flex-grow: 1;
                box-sizing: border-box;
            }
            .input-container {
                display: flex;
                align-items: center;
                justify-content: space-between;
                padding: 10px;
                background-color: #f5f5f5;
                border-top: 1px solid #ccc;
                width: 100%;
                box-sizing: border-box;
            }
            button {
                padding: 10px;
                border: none;
                cursor: pointer;
                background-color: #007bff;
                color: #fff;
                font-size: 16px;
                flex-shrink: 0;
            }
            .user-message {
                background-color: #cce5ff;
                border-radius: 5px;
                align-self: flex-end;
                max-width: 70%;
                margin-left: auto;
                margin-right: 10px;
                margin-bottom: 10px;
                animation: slideInFromRight 0.5s ease;
            }
            .bot-message {
                background-color: #d1ecf1;
                border-radius: 5px;
                align-self: flex-start;
                max-width: 70%;
                margin-bottom: 10px;
                animation: slideInFromLeft 0.5s ease;
            }

            @keyframes fadeIn {
                0% {
                    opacity: 0;
                }
                100% {
                    opacity: 1;
                }
            }

            @keyframes slideInFromRight {
                0% {
                    transform: translateX(100%);
                }
                100% {
                    transform: translateX(0);
                }
            }

            @keyframes slideInFromLeft {
                0% {
                    transform: translateX(-100%);
                }
                100% {
                    transform: translateX(0);
                }
            }
        </style>
    </head>
    <body>
        <div class="container">
            <div class="chat-container">
                <div id="chat-box" class="chat-box"></div>
                <div class="input-container">
                    <input id="chat-input" class="chat-input" type="text" placeholder="Type your message here...">
                    <button id="send-button">Send</button>
                </div>
            </div>
        </div>
        <script>
            const chatBox = document.getElementById('chat-box');
            const chatInput = document.getElementById('chat-input');
            const sendButton = document.getElementById('send-button');

            function appendMessage(text, isUser) {
                const div = document.createElement('div');
                div.className = isUser ? 'user-message' : 'bot-message';
                div.innerHTML = text;
                chatBox.appendChild(div);
                chatBox.scrollTop = chatBox.scrollHeight;
            }

            async function sendMessage() {
                const message = chatInput.value;
                if (message.trim() === '') return;
                appendMessage(message, true);
                chatInput.value = '';

                const response = await fetch('/generate', {
                    method: 'POST',
                    headers: {
                        'Content-Type': 'application/json',
                    },
                    body: JSON.stringify({ query: message }),
                });

                const data = await response.json();
                if (data) {
                    appendMessage(data.responses.join('<br>'), false);

                    if (data.music) {
                        const audio = new Audio('data:audio/wav;base64,' + data.music);
                        audio.play();
                    }

                    if (data.image) {
                        const img = document.createElement('img');
                        img.src = 'data:image/png;base64,' + data.image;
                        chatBox.appendChild(img);
                    }
                }
            }

            sendButton.addEventListener('click', sendMessage);
            chatInput.addEventListener('keypress', (e) => {
                if (e.key === 'Enter') {
                    sendMessage();
                }
            });
        </script>
    </body>
    </html>
    """
    return HTMLResponse(content=html_code)

# Ruta para generar contenido basado en la consulta
@app.post('/generate')
async def generate_content(query: str):
    async def generate_unique_response(q):
        attempts = 0
        while attempts < 5:
            responses = await generate_responses(q)
            unique_responses = [response for response in responses if is_unique(response)]
            if unique_responses:
                parts = []
                for response in unique_responses:
                    parts.extend(split_response(response, model_properties[next(iter(model_dict))]['tokenizer']))
                return parts
            attempts += 1
        raise HTTPException(status_code=500, detail="No unique response found after multiple attempts")

    def is_unique(response):
        if response in UNIQUE_RESPONSES:
            return False
        else:
            UNIQUE_RESPONSES.add(response)
            return True

    async def generate_responses(q):
        responses = []
        for model_name, (model, tokenizer) in model_dict.items():
            input_ids = tokenizer.encode(q, return_tensors='pt')
            output = model.generate(
                input_ids,
                max_length=MAX_TOKENS,
                num_return_sequences=1,
                temperature=TEMPERATURE,
                top_p=TOP_PROBABILITY,
                top_k=TOP_K,
                frequency_penalty=FREQUENCY_PENALTY
            )
            response = tokenizer.decode(output[0], skip_special_tokens=True)
            responses.append(response)
        return responses

    async def train_model():
        global TRAINING_DATA
        if not TRAINING_DATA:
            raise ValueError("No training data available")
        
        model_name = 'gpt2'  # Nombre del modelo a usar para entrenamiento
        tokenizer = AutoTokenizer.from_pretrained(model_name)
        model = AutoModelForCausalLM.from_pretrained(model_name)
        training_args = TrainingArguments(
            output_dir='./results',  # Directorio de resultados (en memoria, se puede ignorar)
            per_device_train_batch_size=4,
            num_train_epochs=1,
            save_steps=10_000,
            save_total_limit=2,
        )
        trainer = Trainer(
            model=model,
            args=training_args,
            train_dataset=TRAINING_DATA
        )
        trainer.train()
        # Guardar el modelo entrenado en Redis
        model_key = "model:trained"
        tokenizer_key = "tokenizer:trained"
        store_in_redis(model_key, torch.save(model, io.BytesIO()))
        store_in_redis(tokenizer_key, torch.save(tokenizer, io.BytesIO()))
        return model, tokenizer

    async def auto_learn():
        global TRAINING_DATA
        if message_history:
            new_data = "\n".join(message_history)
            TRAINING_DATA.append(new_data)
            await train_model()

    async def auto_learn_music():
        global MUSIC_TRAINING_DATA
        if MUSIC_TRAINING_DATA:
            inputs = musicgen_pipeline.tokenizer(MUSIC_TRAINING_DATA, return_tensors="pt", padding=True)
            model = musicgen_pipeline.model
            model.train()
            optimizer = torch.optim.Adam(model.parameters(), lr=5e-5)
            loss_fn = torch.nn.CrossEntropyLoss()

            for epoch in range(1):
                outputs = model(**inputs)
                loss = loss_fn(outputs.logits, inputs['labels'])
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()

            MUSIC_TRAINING_DATA = []

    async def auto_learn_images():
        global IMAGE_TRAINING_DATA
        if IMAGE_TRAINING_DATA:
            for image_data in IMAGE_TRAINING_DATA:
                image = Image.open(io.BytesIO(image_data))
                image_tensor = torch.tensor(np.array(image)).unsqueeze(0)  # Adaptar según el modelo
                # Implementar el entrenamiento del modelo aquí con `image_tensor`
                model = image_pipeline.model
                model.train()
                optimizer = torch.optim.Adam(model.parameters(), lr=1e-5)
                loss_fn = torch.nn.MSELoss()
                target_tensor = torch.zeros_like(image_tensor)  # Definir `target_tensor` según sea necesario
                for epoch in range(1):
                    outputs = model(image_tensor)
                    loss = loss_fn(outputs, target_tensor)
                    optimizer.zero_grad()
                    loss.backward()
                    optimizer.step()
            IMAGE_TRAINING_DATA = []

    def generate_music_from_api(prompt: str) -> bytes:
        # Llamada a la API para generar música
        audio = generate_music(prompt)
        store_in_redis(f"music:{prompt}", audio)  # Almacenar música en Redis
        return audio

    def generate_image_from_api(prompt: str) -> bytes:
        # Llamada a la API para generar imágenes
        image = generate_image(prompt)
        store_in_redis(f"image:{prompt}", image)  # Almacenar imagen en Redis
        return image

    try:
        tokenizer = model_properties[next(iter(model_dict))]['tokenizer']
        final_responses = await generate_unique_response(query)
        await auto_learn()
        await auto_learn_music()
        await auto_learn_images()
        music = generate_music_from_api(query)
        image = generate_image_from_api(query)
        
        # Convertir la imagen a base64 para mostrarla en la interfaz web
        buffered = io.BytesIO(image)
        img_str = base64.b64encode(buffered.getvalue()).decode('utf-8')
        
        return {"responses": final_responses, "music": base64.b64encode(music).decode('utf-8'), "image": img_str}
    except Exception as e:
        logger.error(f"Error processing the request: {e}")
        raise HTTPException(status_code=500, detail="Error processing the request")

# Ruta para generación de música
@app.post('/music')
async def generate_music_endpoint(prompt: str):
    try:
        music = generate_music_from_api(prompt)
        return {"music": base64.b64encode(music).decode('utf-8')}
    except Exception as e:
        logger.error(f"Error generating music: {e}")
        raise HTTPException(status_code=500, detail="Error generating music")

# Ruta para generación de imágenes
@app.post('/image')
async def generate_image_endpoint(prompt: str):
    try:
        image = generate_image_from_api(prompt)
        img_str = base64.b64encode(image).decode('utf-8')
        return {"image": img_str}
    except Exception as e:
        logger.error(f"Error generating image: {e}")
        raise HTTPException(status_code=500, detail="Error generating image")

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)