app_wan.py

# demo_aduc_wan.py
import gradio as gr
import torch
import logging
import os
from PIL import Image
from typing import List, Dict, Any

# --- Importa os Pilares da nossa Arquitetatura ADUC ---
from aduc_framework.managers.vae_wan_manager import vae_wan_manager_singleton as VaeWan
from aduc_framework.managers.wan_manager import wan_manager_singleton as Wan
from aduc_framework.types import LatentConditioningItem
from aduc_framework.tools.video_encode_tool import video_encode_tool_singleton as VideoTool

# --- Configuração do Logging ---
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - [%(name)s] - %(message)s')
logger = logging.getLogger("ADUC_WAN_DEMO")

# --- Lógica Principal da Geração (Adaptada para Gradio) ---

def generate_aduc_video(
    image1: Image.Image, strength1: float, frame_idx1: int,
    image2: Image.Image, strength2: float, frame_idx2: int,
    image3: Image.Image, strength3: float, frame_idx3: int,
    motion_prompt: str,
    progress=gr.Progress(track_tqdm=True)
):
    """
    Função principal que orquestra a geração de vídeo e é chamada pela interface Gradio.
    """
    try:
        # --- 1. Validação e Montagem do Plano de Produção ---
        progress(0, desc="[Diretor] Validando o plano de produção...")
        
        production_plan = []
        if image1: production_plan.append({"image": image1, "frame_index": int(frame_idx1), "strength": strength1})
        if image2: production_plan.append({"image": image2, "frame_index": int(frame_idx2), "strength": strength2})
        if image3: production_plan.append({"image": image3, "frame_index": int(frame_idx3), "strength": strength3})
        
        if not production_plan:
            raise gr.Error("Por favor, forneça pelo menos uma imagem para iniciar a geração.")

        resolution = (480, 832) # Resolução padrão para o teste

        # --- 2. Pré-Produção: Conversão para Latentes ---
        progress(0.1, desc="[VaeWan] Convertendo imagens para o espaço latente...")
        pil_images = [item["image"] for item in production_plan]
        latent_tensors = VaeWan.encode_batch(pil_images, target_resolution=resolution)
        
        # --- 3. Estruturação: Montando as Condições ADUC ---
        progress(0.2, desc="[Diretor] Montando ordens de serviço (LatentConditioningItems)...")
        conditioning_items = [
            LatentConditioningItem(
                latent_tensor=latent_tensors[i],
                media_frame_number=item["frame_index"],
                conditioning_strength=item["strength"]
            ) for i, item in enumerate(production_plan)
        ]

        # --- 4. Produção: Geração pelo WanManager ---
        progress(0.3, desc="[WanManager] Luz, Câmera, Ação! Gerando fragmento latente (pode levar um momento)...")
        job_params = {
            "conditioning_items_data": conditioning_items,
            "motion_prompt": motion_prompt,
            "height": resolution[0],
            "width": resolution[1],
            "video_total_frames": 81,
            "num_inference_steps": 8,
        }
        video_latents, _ = Wan.generate_latent_fragment(**job_params)

        # --- 5. Pós-Produção: Decodificação ---
        progress(0.8, desc="[VaeWan] Decodificando o resultado final para pixels...")
        video_pixels = VaeWan.decode(video_latents)

        # --- 6. Salvamento ---
        progress(0.9, desc="[VideoTool] Salvando o arquivo de vídeo final...")
        output_path = "demo_aduc_wan_output.mp4"
        VideoTool.save_video_from_tensor(video_pixels, path=output_path, fps=16)

        progress(1.0, desc="Produção Concluída!")
        return output_path

    except Exception as e:
        logger.error("Ocorreu um erro durante a geração do vídeo.", exc_info=True)
        # Lança um erro que será exibido na interface do Gradio
        raise gr.Error(f"Falha na geração: {e}")


# --- Construção da Interface Gráfica (UI) ---

with gr.Blocks(theme=gr.themes.Soft(), title="ADUC-SDR Wan Demo") as demo:
    gr.Markdown(
        """
        # 🎬 Demo ADUC-SDR com WanManager (Lightning)
        ### Bem-vindo, Mestre Deformes!
        Esta interface permite controlar o especialista `WanManager` usando a metodologia ADUC.
        Faça o upload de até 3 keyframes, defina suas posições e forças no vídeo e escreva um prompt para guiar a animação.
        """
    )

    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### Keyframe 1 (Inicial)")
            img1 = gr.Image(type="pil", label="Imagem de Início")
            strength1 = gr.Slider(minimum=0.0, maximum=1.0, value=1.0, label="Força (Strength)", info="1.0 = o vídeo começa EXATAMENTE nesta imagem.")
            frame_idx1 = gr.Number(value=0, label="Índice do Frame", info="Em que ponto do vídeo este frame aparece (0 é o início).", precision=0)

        with gr.Column(scale=1):
            gr.Markdown("### Keyframe 2 (Intermediário)")
            img2 = gr.Image(type="pil", label="Imagem do Meio")
            strength2 = gr.Slider(minimum=0.0, maximum=1.0, value=0.6, label="Força (Strength)", info="< 1.0 = o vídeo é INFLUENCIADO por esta imagem.")
            frame_idx2 = gr.Number(value=40, label="Índice do Frame", info="O padrão é o meio do vídeo (frame 40 de 81).", precision=0)

        with gr.Column(scale=1):
            gr.Markdown("### Keyframe 3 (Final)")
            img3 = gr.Image(type="pil", label="Imagem Final")
            strength3 = gr.Slider(minimum=0.0, maximum=1.0, value=1.0, label="Força (Strength)", info="1.0 = o vídeo termina EXATAMENTE nesta imagem.")
            frame_idx3 = gr.Number(value=80, label="Índice do Frame", info="O padrão é o último frame do vídeo.", precision=0)

    with gr.Row():
        motion_prompt_input = gr.Textbox(
            label="📝 Prompt de Movimento",
            placeholder="Ex: a beautiful cinematic shot of a majestic landscape changing seasons",
            lines=2,
            scale=3
        )
        generate_btn = gr.Button("🎬 Gerar Vídeo", variant="primary", scale=1)

    with gr.Row():
        output_video = gr.Video(label="Resultado da Produção", interactive=False)

    # Conecta o botão à função de geração
    generate_btn.click(
        fn=generate_aduc_video,
        inputs=[
            img1, strength1, frame_idx1,
            img2, strength2, frame_idx2,
            img3, strength3, frame_idx3,
            motion_prompt_input
        ],
        outputs=[output_video]
    )
    
    gr.Examples(
        examples=[
            ["examples/frame_1.png", 1.0, 0, "examples/frame_2.png", 0.7, 40, "examples/frame_3.png", 1.0, 80, "A slow zoom out revealing a vast alien desert at sunset"],
        ],
        inputs=[img1, strength1, frame_idx1, img2, strength2, frame_idx2, img3, strength3, frame_idx3, motion_prompt_input],
        outputs=[output_video],
        fn=generate_aduc_video,
        cache_examples=False, # Desativa o cache para sempre rodar a geração
    )

if __name__ == "__main__":
    # Cria a pasta e as imagens de exemplo se não existirem
    if not os.path.exists("examples"):
        os.makedirs("examples")
    try:
        Image.new('RGB', (832, 480), color = (73, 109, 137)).save("examples/frame_1.png")
        Image.new('RGB', (832, 480), color = (173, 109, 237)).save("examples/frame_2.png")
        Image.new('RGB', (832, 480), color = (255, 255, 0)).save("examples/frame_3.png")
    except:
        pass # Evita que o app quebre se não tiver permissão de escrita

    os.makedirs(WORKSPACE_DIR, exist_ok=True)
    logger.info("Aplicação Gradio pronta. Lançando interface...")
    demo.launch(
        debug=True,
        server_name="0.0.0.0",
        server_port=int(os.getenv("PORT", "7860")),
    )