app_simple.py

# IMPORTANTE: Importar spaces ANTES de torch para ZeroGPU
import spaces  # Para usar ZeroGPU H200

import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
from diffusers import StableDiffusionPipeline, DiffusionPipeline
import requests
from PIL import Image
import io
import base64
import os
import time
import numpy as np
import random
from huggingface_hub import login
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel

print("🚀 Iniciando NTIA Space con ZeroGPU H200...")
print(f"📁 Directorio actual: {os.getcwd()}")
print(f"🐍 Python version: {os.sys.version}")

# Optimización para ZeroGPU H200
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"🖥️ Dispositivo detectado: {device}")
print(f"🔥 CUDA disponible: {torch.cuda.is_available()}")

if torch.cuda.is_available():
    print(f"🎮 GPU: {torch.cuda.get_device_name(0)}")
    print(f"💾 Memoria GPU: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.1f} GB")
    print("🚀 ZeroGPU H200 detectado - Optimizando para máximo rendimiento")
    
    # Configuración optimizada para H200
    torch_dtype = torch.float16  # Usar float16 para mayor velocidad
    print("⚡ Usando torch.float16 para H200")
    
    # Optimizaciones adicionales para H200
    torch.backends.cudnn.benchmark = True
    torch.backends.cuda.matmul.allow_tf32 = True
    torch.backends.cudnn.allow_tf32 = True
    print("🔧 Optimizaciones CUDA habilitadas para H200")
else:
    torch_dtype = torch.float32
    print("🐌 Usando torch.float32 para CPU")

# Configurar autenticación con Hugging Face
HF_TOKEN = os.getenv("HF_TOKEN") or os.getenv("HUGGING_FACE_HUB_TOKEN")
if HF_TOKEN:
    try:
        print(f"🔑 Token detectado: {HF_TOKEN[:10]}...")
        login(token=HF_TOKEN)
        print("✅ Autenticado con Hugging Face")
    except Exception as e:
        print(f"⚠️ Error de autenticación: {e}")
else:
    print("⚠️ No se encontró HF_TOKEN - modelos gated no estarán disponibles")

# Configuración de modelos libres
MODELS = {
    "text": {
        "microsoft/DialoGPT-medium": "Chat conversacional",
        "gpt2": "Generación de texto",
        "distilgpt2": "GPT-2 optimizado"
    },
    "image": {
        "CompVis/stable-diffusion-v1-4": "Stable Diffusion v1.4 (Libre)",
        "stabilityai/stable-diffusion-2-1": "Stable Diffusion 2.1",
        "stabilityai/stable-diffusion-xl-base-1.0": "SDXL Base",
        "prompthero/openjourney": "Midjourney Style",
        "hakurei/waifu-diffusion": "Waifu Diffusion"
    },
    "video": {
        "damo-vilab/text-to-video-ms-1.7b": "Text-to-Video MS 1.7B (Libre)",
        "cerspense/zeroscope_v2_576w": "Zeroscope v2 576w (Libre)"
    }
}

# Cache para los modelos
model_cache = {}

def load_text_model(model_name):
    """Cargar modelo de texto"""
    if model_name not in model_cache:
        print(f"Cargando modelo de texto: {model_name}")
        
        try:
            tokenizer = AutoTokenizer.from_pretrained(model_name)
            model = AutoModelForCausalLM.from_pretrained(model_name)
            
            if "dialogpt" in model_name.lower():
                tokenizer.pad_token = tokenizer.eos_token
                model.config.pad_token_id = model.config.eos_token_id
            
            model_cache[model_name] = {
                "tokenizer": tokenizer,
                "model": model,
                "type": "text"
            }
            
        except Exception as e:
            print(f"Error cargando modelo de texto {model_name}: {e}")
            # Fallback
            tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-medium")
            model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-medium")
            tokenizer.pad_token = tokenizer.eos_token
            model.config.pad_token_id = model.config.eos_token_id
            
            model_cache[model_name] = {
                "tokenizer": tokenizer,
                "model": model,
                "type": "text"
            }
    
    return model_cache[model_name]

def load_image_model(model_name):
    """Cargar modelo de imagen optimizado para H200"""
    if model_name not in model_cache:
        print(f"Cargando modelo de imagen: {model_name}")
        
        try:
            pipe = StableDiffusionPipeline.from_pretrained(
                model_name,
                torch_dtype=torch_dtype,
                safety_checker=None
            )
            
            pipe = pipe.to(device)
            
            # Optimizaciones para H200
            if torch.cuda.is_available():
                pipe.enable_attention_slicing()
                pipe.enable_vae_slicing()
                
                if hasattr(pipe, 'enable_xformers_memory_efficient_attention'):
                    try:
                        pipe.enable_xformers_memory_efficient_attention()
                    except:
                        pass
            
            model_cache[model_name] = pipe
            
        except Exception as e:
            print(f"Error cargando modelo {model_name}: {e}")
            # Fallback
            pipe = StableDiffusionPipeline.from_pretrained(
                "CompVis/stable-diffusion-v1-4",
                torch_dtype=torch_dtype,
                safety_checker=None
            )
            pipe = pipe.to(device)
            model_cache[model_name] = pipe
    
    return model_cache[model_name]

def generate_text(prompt, model_name, max_length=100):
    """Generar texto con el modelo seleccionado"""
    try:
        model_data = load_text_model(model_name)
        tokenizer = model_data["tokenizer"]
        model = model_data["model"]
        
        inputs = tokenizer.encode(prompt, return_tensors="pt")
        
        with torch.no_grad():
            outputs = model.generate(
                inputs,
                max_length=max_length,
                num_return_sequences=1,
                temperature=0.7,
                do_sample=True,
                pad_token_id=tokenizer.eos_token_id
            )
        
        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
        
        if "dialogpt" in model_name.lower():
            response = response.replace(prompt, "").strip()
        
        return response
        
    except Exception as e:
        return f"Error generando texto: {str(e)}"

@spaces.GPU
def generate_image(prompt, model_name, negative_prompt="", seed=0, width=1024, height=1024, guidance_scale=7.5, num_inference_steps=20):
    """Generar imagen optimizada para H200"""
    try:
        print(f"Generando imagen: {prompt}")
        
        pipe = load_image_model(model_name)
        
        generator = torch.Generator(device=device).manual_seed(seed)
        
        result = pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,
            height=height,
            width=width,
            guidance_scale=guidance_scale,
            num_inference_steps=num_inference_steps,
            generator=generator
        )
        
        image = result.images[0]
        print("✅ Imagen generada exitosamente")
        return image
        
    except Exception as e:
        print(f"Error generando imagen: {e}")
        error_image = Image.new('RGB', (512, 512), color='red')
        return error_image

def chat_with_model(message, history, model_name):
    """Función de chat para DialoGPT"""
    try:
        model_data = load_text_model(model_name)
        tokenizer = model_data["tokenizer"]
        model = model_data["model"]
        
        conversation = ""
        for msg in history:
            if msg["role"] == "user":
                conversation += f"User: {msg['content']}\n"
            elif msg["role"] == "assistant":
                conversation += f"Assistant: {msg['content']}\n"
        
        conversation += f"User: {message}\nAssistant:"
        
        inputs = tokenizer.encode(conversation, return_tensors="pt", truncation=True, max_length=512)
        
        with torch.no_grad():
            outputs = model.generate(
                inputs,
                max_length=inputs.shape[1] + 50,
                temperature=0.7,
                do_sample=True,
                pad_token_id=tokenizer.eos_token_id
            )
        
        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
        response = response.split("Assistant:")[-1].strip()
        
        history.append({"role": "user", "content": message})
        history.append({"role": "assistant", "content": response})
        
        return history
        
    except Exception as e:
        error_msg = f"Error en el chat: {str(e)}"
        history.append({"role": "user", "content": message})
        history.append({"role": "assistant", "content": error_msg})
        return history

# Interfaz de Gradio
with gr.Blocks(title="Modelos Libres de IA", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🤖 Modelos Libres de IA")
    gr.Markdown("### Genera texto e imágenes sin límites de cuota")
    
    with gr.Tabs():
        # Tab de Generación de Texto
        with gr.TabItem("📝 Generación de Texto"):
            with gr.Row():
                with gr.Column():
                    text_model = gr.Dropdown(
                        choices=list(MODELS["text"].keys()),
                        value="microsoft/DialoGPT-medium",
                        label="Modelo de Texto"
                    )
                    text_prompt = gr.Textbox(
                        label="Prompt",
                        placeholder="Escribe tu prompt aquí...",
                        lines=3
                    )
                    max_length = gr.Slider(
                        minimum=50,
                        maximum=200,
                        value=100,
                        step=10,
                        label="Longitud máxima"
                    )
                    text_btn = gr.Button("Generar Texto", variant="primary")
                
                with gr.Column():
                    text_output = gr.Textbox(
                        label="Resultado",
                        lines=10,
                        interactive=False
                    )
            
            text_btn.click(
                generate_text,
                inputs=[text_prompt, text_model, max_length],
                outputs=text_output
            )
        
        # Tab de Chat
        with gr.TabItem("💬 Chat"):
            with gr.Row():
                with gr.Column():
                    chat_model = gr.Dropdown(
                        choices=list(MODELS["text"].keys()),
                        value="microsoft/DialoGPT-medium",
                        label="Modelo de Chat"
                    )
                
                with gr.Column():
                    chatbot = gr.Chatbot(
                        label="Chat",
                        height=400,
                        type="messages"
                    )
                    chat_input = gr.Textbox(
                        label="Mensaje",
                        placeholder="Escribe tu mensaje...",
                        lines=2
                    )
                    chat_btn = gr.Button("Enviar", variant="primary")
            
            chat_btn.click(
                chat_with_model,
                inputs=[chat_input, chatbot, chat_model],
                outputs=[chatbot]
            )
            
            chat_input.submit(
                chat_with_model,
                inputs=[chat_input, chatbot, chat_model],
                outputs=[chatbot]
            )
        
        # Tab de Generación de Imágenes
        with gr.TabItem("🎨 Generación de Imágenes"):
            with gr.Row():
                with gr.Column():
                    image_model = gr.Dropdown(
                        choices=list(MODELS["image"].keys()),
                        value="CompVis/stable-diffusion-v1-4",
                        label="Modelo"
                    )
                    
                    image_prompt = gr.Textbox(
                        label="Prompt",
                        placeholder="Describe la imagen que quieres generar...",
                        lines=3
                    )
                    
                    negative_prompt = gr.Textbox(
                        label="Negative prompt",
                        placeholder="Enter a negative prompt (optional)",
                        lines=2
                    )
                    
                    with gr.Row():
                        seed = gr.Slider(
                            minimum=0,
                            maximum=2147483647,
                            value=324354329,
                            step=1,
                            label="Seed"
                        )
                        guidance_scale = gr.Slider(
                            minimum=0,
                            maximum=20,
                            value=7.5,
                            step=0.1,
                            label="Guidance scale"
                        )
                    
                    with gr.Row():
                        width = gr.Slider(
                            minimum=256,
                            maximum=1024,
                            value=1024,
                            step=64,
                            label="Width"
                        )
                        height = gr.Slider(
                            minimum=256,
                            maximum=1024,
                            value=1024,
                            step=64,
                            label="Height"
                        )
                    
                    num_inference_steps = gr.Slider(
                        minimum=1,
                        maximum=100,
                        value=20,
                        step=1,
                        label="Number of inference steps"
                    )
                    
                    image_btn = gr.Button("Generar Imagen", variant="primary")
                
                with gr.Column():
                    examples = gr.Examples(
                        examples=[
                            ["Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"],
                            ["An astronaut riding a green horse"],
                            ["A delicious ceviche cheesecake slice"],
                            ["Futuristic AI assistant in a glowing galaxy, neon lights, sci-fi style, cinematic"]
                        ],
                        inputs=image_prompt
                    )
                    
                    image_output = gr.Image(
                        label="Imagen Generada",
                        type="pil"
                    )
            
            image_btn.click(
                generate_image,
                inputs=[
                    image_prompt, 
                    image_model, 
                    negative_prompt,
                    seed,
                    width,
                    height,
                    guidance_scale,
                    num_inference_steps
                ],
                outputs=image_output
            )

# Configuración para Hugging Face Spaces
if __name__ == "__main__":
    demo.launch()