playground/test/list_models.py

import json
import os
import base64
import mimetypes
import time
from datetime import datetime
from pathlib import Path
from google import genai
from google.genai import types
from google.genai.types import GenerateVideosConfig

# Leer credenciales
credentials_path = Path(__file__).parent.parent / "bnt-ia-innovacion-new.json"

with open(credentials_path, 'r') as f:
    credentials_json = json.load(f)

# Configurar variable de entorno con las credenciales
os.environ['GOOGLE_APPLICATION_CREDENTIALS'] = str(credentials_path)

# Crear cliente de genai
client = genai.Client(
    vertexai=True,
    project=credentials_json.get('project_id'),
    location='us-central1'
)

# Crear directorio para videos
output_dir = Path("generated_videos")
output_dir.mkdir(exist_ok=True)

print("="*60)
print("CHAT DE GENERACIÓN DE VIDEOS CON GOOGLE VEO")
print("="*60)
print(f"\nProyecto: {credentials_json.get('project_id')}")
print(f"Ubicación: us-central1")
print(f"Modelo: veo-002")
print(f"Directorio de salida: {output_dir}")
print("="*60)


def save_video(video_data, prompt, index=0):
    """Guarda el video generado en disco"""
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    # Sanitizar el prompt para usar como parte del nombre
    safe_prompt = "".join(c for c in prompt[:30] if c.isalnum() or c in (' ', '_')).rstrip()
    safe_prompt = safe_prompt.replace(' ', '_')

    filename = f"{timestamp}_{safe_prompt}_{index}.mp4"
    filepath = output_dir / filename

    # Decodificar y guardar el video
    if isinstance(video_data, str):
        # Si es base64
        video_bytes = base64.b64decode(video_data)
    else:
        video_bytes = video_data

    with open(filepath, 'wb') as f:
        f.write(video_bytes)

    return filepath


def load_reference_media(media_path):
    """Carga una imagen o video de referencia y la convierte a Part"""
    try:
        media_path = Path(media_path)

        if not media_path.exists():
            print(f"✗ Archivo no encontrado: {media_path}")
            return None

        # Detectar tipo MIME
        mime_type, _ = mimetypes.guess_type(str(media_path))
        if mime_type is None:
            print(f"✗ No se pudo detectar el tipo de archivo: {media_path}")
            return None

        if not (mime_type.startswith('image/') or mime_type.startswith('video/')):
            print(f"✗ Archivo no es una imagen o video válido: {media_path}")
            return None

        # Leer archivo
        with open(media_path, 'rb') as f:
            media_bytes = f.read()

        # Crear Part con el archivo
        return types.Part.from_bytes(
            data=media_bytes,
            mime_type=mime_type
        )

    except Exception as e:
        print(f"✗ Error al cargar archivo: {e}")
        return None


def generate_video(prompt, reference_media=None, duration=5):
    """Genera un video y lo guarda

    Args:
        prompt: Texto descriptivo para generar el video
        reference_media: Lista de rutas a imágenes de referencia (opcional)
        duration: Duración del video en segundos (5 u 8)
    """
    print(f"\n{'='*60}")
    print(f"Generando video: {prompt}")
    print(f"Duración: {duration} segundos")
    if reference_media:
        print(f"Con {len(reference_media)} imagen(es) de referencia")
    print(f"{'='*60}\n")

    try:
        # Configurar bucket de salida (temporal en GCS)
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        safe_prompt = "".join(c for c in prompt[:20] if c.isalnum() or c in (' ', '_')).rstrip()
        safe_prompt = safe_prompt.replace(' ', '_')
        output_gcs_uri = f"gs://bnt-ia-innovacion-videos/veo_output/{timestamp}_{safe_prompt}/"

        # Preparar imagen de referencia si existe
        reference_image = None
        if reference_media and len(reference_media) > 0:
            # Para Veo, solo se puede usar una imagen de referencia
            img_path = reference_media[0]
            print(f"📸 Cargando imagen de referencia: {img_path}")

            # Subir imagen a GCS o usar directamente si ya está en GCS
            if img_path.startswith('gs://'):
                reference_image = types.Image(gcs_uri=img_path)
                print(f"✓ Usando imagen de GCS: {img_path}")
            else:
                # Por ahora mostrar advertencia
                print(f"⚠ La imagen debe estar en Google Cloud Storage (gs://)")
                print(f"⚠ Continuando sin imagen de referencia...")
                reference_image = None

        print("\n⏳ Generando video con Veo (esto puede tardar varios minutos)...")
        print(f"📂 El video se guardará en: {output_gcs_uri}")
        start_time = time.time()

        # Configuración de Veo
        config = GenerateVideosConfig(
            aspect_ratio="16:9",
            output_gcs_uri=output_gcs_uri,
        )

        # Generar video con Veo
        if reference_image:
            operation = client.models.generate_videos(
                model="veo-3.1-generate-001",
                prompt=prompt,
                reference_image=reference_image,
                config=config,
            )
        else:
            operation = client.models.generate_videos(
                model="veo-3.1-generate-001",
                prompt=prompt,
                config=config,
            )

        print(f"🔄 Operación iniciada: {operation.name}")
        print("⏳ Esperando a que se complete la generación...")

        # Polling de la operación cada 15 segundos
        poll_count = 0
        while not operation.done:
            time.sleep(15)
            operation = client.operations.get(operation)
            poll_count += 1
            elapsed = time.time() - start_time
            print(f"⏱ {elapsed:.0f}s - Verificación #{poll_count} - Estado: {'Completado' if operation.done else 'En proceso...'}")

        elapsed_time = time.time() - start_time
        print(f"\n✓ Generación completada en {elapsed_time:.2f} segundos")

        # Obtener el resultado
        if hasattr(operation, 'result') and operation.result:
            if hasattr(operation.result, 'generated_videos') and operation.result.generated_videos:
                video_uri = operation.result.generated_videos[0].video.uri
                print(f"\n✓ Video generado exitosamente!")
                print(f"📹 URI del video en GCS: {video_uri}")

                # Intentar descargar el video de GCS
                try:
                    from google.cloud import storage

                    # Parsear la URI de GCS
                    gcs_path = video_uri.replace('gs://', '')
                    bucket_name = gcs_path.split('/')[0]
                    blob_name = '/'.join(gcs_path.split('/')[1:])

                    print(f"\n⬇ Descargando video de GCS...")
                    storage_client = storage.Client(project=credentials_json.get('project_id'))
                    bucket = storage_client.bucket(bucket_name)
                    blob = bucket.blob(blob_name)

                    # Guardar localmente
                    local_filename = f"{timestamp}_{safe_prompt}.mp4"
                    local_filepath = output_dir / local_filename
                    blob.download_to_filename(str(local_filepath))

                    print(f"✓ Video descargado en: {local_filepath}")
                    return True

                except ImportError:
                    print("\n⚠ Para descargar automáticamente, instala: pip install google-cloud-storage")
                    print(f"📹 Puedes descargar el video manualmente desde: {video_uri}")
                    return True
                except Exception as e:
                    print(f"\n⚠ No se pudo descargar automáticamente: {e}")
                    print(f"📹 Puedes descargar el video manualmente desde: {video_uri}")
                    return True
            else:
                print("\n⚠ No se generaron videos en la respuesta")
                print(f"DEBUG - Resultado: {operation.result}")
                return False
        else:
            print("\n⚠ La operación no tiene resultado")
            print(f"DEBUG - Operación: {operation}")
            return False

    except Exception as e:
        print(f"\n✗ Error al generar video: {str(e)}")
        import traceback
        print("\nTraceback completo:")
        traceback.print_exc()
        return False


# Chat interactivo
print("\nComandos:")
print("  - Escribe un prompt para generar un video")
print("  - Para especificar duración: prompt [5s] o [8s]")
print("  - Para usar archivos de referencia: prompt | ruta1.jpg | video.mp4")
print("  - Ejemplo: un gato saltando en slow motion [8s] | referencia.jpg")
print("  - 'salir' para terminar\n")

while True:
    try:
        user_input = input("\n🎬 Prompt: ").strip()

        if not user_input:
            continue

        if user_input.lower() == 'salir':
            print("\n👋 ¡Hasta pronto!")
            break

        # Extraer duración si está especificada
        duration = 5  # Por defecto 5 segundos
        if '[5s]' in user_input:
            duration = 5
            user_input = user_input.replace('[5s]', '').strip()
        elif '[8s]' in user_input:
            duration = 8
            user_input = user_input.replace('[8s]', '').strip()

        # Verificar si hay archivos de referencia
        if '|' in user_input:
            parts = [p.strip() for p in user_input.split('|')]
            prompt = parts[0]
            reference_media = parts[1:] if len(parts) > 1 else None
            generate_video(prompt, reference_media, duration)
        else:
            generate_video(user_input, duration=duration)

    except KeyboardInterrupt:
        print("\n\n👋 ¡Hasta pronto!")
        break
    except Exception as e:
        print(f"\n✗ Error: {str(e)}")