app_vince.py

#!/usr/bin/env python3
"""
VINCIE Service UI (Gradio) - Versão Final com Suporte a Multi-GPU

- Interface com 3 abas: Edição Multi-Turno, Texto-para-Vídeo e Composição Multi-Conceito.
- Controles avançados na UI, incluindo seleção de número de GPUs e tamanho do lote (batch size).
- Configuração automática no carregamento da aplicação.
- Projetado para hardware de ponta como 8x L40S.
- Modelo funcional de referência: ByteDance-Seed/VINCIE.
- Interface desenvolvida por Carlex (carlex22@gmail.com).
"""

import os
from pathlib import Path
from typing import List, Tuple, Optional
import gradio as gr

# Adapte este import para o layout do seu projeto, se necessário.
from services.vincie import VincieService

# --- Instanciação do Serviço e Constantes ---
svc = VincieService()
DEFAULT_NEGATIVE_PROMPT = "Worst quality, Normal quality, Low quality, Low res, Blurry, Jpeg artifacts, Grainy, text, logo, watermark, banner, extra digits, signature, subtitling, Bad anatomy, Bad proportions, Deformed, Disconnected limbs, Disfigured, Extra arms, Extra limbs, Extra hands, Fused fingers, Gross proportions, Long neck, Malformed limbs, Mutated, Mutated hands, Mutated limbs, Missing arms, Missing fingers, Poorly drawn hands, Poorly drawn face, Nsfw, Uncensored, Cleavage, Nude, Nipples, Overexposed, Plain background, Grainy, Underexposed, Deformed structures"

# --- Funções Helper ---
def setup_auto() -> str:
    """
    Executa uma configuração idempotente no carregamento da interface.
    Retorna uma string de status para a UI.
    """
    try:
        svc.ensure_repo()
        svc.ensure_model()
        return "Configuração concluída com sucesso: repositório e checkpoint estão prontos."
    except Exception as e:
        # Fornece um feedback de erro mais detalhado para depuração
        import traceback
        print(traceback.format_exc())
        return f"A configuração encontrou um erro: {e}"

def _list_media(out_dir: Path, max_images: int = 24) -> Tuple[List[str], Optional[str]]:
    """
    Enumera as imagens resultantes e o vídeo mais recente de um diretório de saída.
    """
    img_globs = ("*.png", "*.jpg", "*.jpeg", "*.webp")
    # Usa rglob para encontrar imagens em subdiretórios e ordena por tempo de modificação
    images = sorted(
        [p for pat in img_globs for p in out_dir.rglob(pat)],
        key=lambda p: p.stat().st_mtime
    )
    image_paths = [str(p) for p in images[-max_images:]]
    videos = sorted(out_dir.rglob("*.mp4"), key=lambda p: p.stat().st_mtime)
    video_path = str(videos[-1]) if videos else None
    return image_paths, video_path

# --- Funções Handler da UI (com todos os parâmetros) ---
def ui_multi_turn(input_image, turns_text, negative_prompt, seed, steps, cfg_scale, resolution, use_vae_slicing, num_gpus, batch_size):
    """Handler para a aba de Edição Multi-Turno."""
    if not input_image:
        return [], None, "Por favor, forneça uma imagem de entrada."
    if not turns_text or not turns_text.strip():
        return [], None, "Por favor, forneça as instruções de edição (uma por linha)."

    turns = [ln.strip() for ln in turns_text.splitlines() if ln.strip()]
    try:
        out_dir = svc.multi_turn_edit(
            input_image, turns,
            negative_prompt=negative_prompt, seed=int(seed), steps=int(steps), 
            cfg_scale=float(cfg_scale), resolution=int(resolution), use_vae_slicing=use_vae_slicing,
            num_gpus=int(num_gpus), batch_size=int(batch_size)
        )
        imgs, vid = _list_media(Path(out_dir))
        return imgs, vid, f"Saídas salvas em: {out_dir}"
    except Exception as e:
        import traceback
        print(traceback.format_exc())
        return [], None, f"Erro na geração: {e}"

def ui_text_to_video(input_image, prompt, negative_prompt, seed, steps, cfg_scale, resolution, fps, use_vae_slicing, num_gpus, batch_size):
    """Handler para a aba de Texto-para-Vídeo."""
    if not input_image:
        return None, "Por favor, forneça uma imagem de entrada (frame inicial)."
    if not prompt or not prompt.strip():
        return None, "Por favor, forneça um prompt para o vídeo."

    try:
        out_dir = svc.text_to_video(
            input_image, prompt,
            negative_prompt=negative_prompt, seed=int(seed), steps=int(steps), 
            cfg_scale=float(cfg_scale), resolution=int(resolution), fps=int(fps), use_vae_slicing=use_vae_slicing,
            num_gpus=int(num_gpus), batch_size=int(batch_size)
        )
        _, vid = _list_media(Path(out_dir))
        return vid, f"Vídeo salvo em: {out_dir}"
    except Exception as e:
        import traceback
        print(traceback.format_exc())
        return None, f"Erro na geração: {e}"

def ui_multi_concept(files, descs_text, final_prompt):
    """Handler para a aba de Composição Multi-Conceito."""
    if not files: return [], None, "Por favor, faça o upload das imagens de conceito."
    if not descs_text: return [], None, "Por favor, forneça as descrições (uma por linha)."
    if not final_prompt: return [], None, "Por favor, forneça um prompt final."
    
    descs = [ln.strip() for ln in descs_text.splitlines() if ln.strip()]
    if len(descs) != len(files): return [], None, f"O número de descrições ({len(descs)}) deve ser igual ao de imagens ({len(files)})."

    try:
        out_dir = svc.multi_concept_compose(files, descs, final_prompt)
        imgs, vid = _list_media(Path(out_dir))
        return imgs, vid, f"Saídas salvas em: {out_dir}"
    except Exception as e:
        import traceback
        print(traceback.format_exc())
        return [], None, f"Erro na geração: {e}"

# --- Definição da Interface Gradio Completa ---
with gr.Blocks(title="VINCIE Service", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# VINCIE Service — Geração Distribuída com Controles Avançados")
    gr.Markdown("- **Interface por:** Carlex (carlex22@gmail.com | GitHub: carlex22)")
    with gr.Row():
        setup_out = gr.Textbox(label="Status da Configuração", interactive=False)

    with gr.Tab("Edição Multi-Turno"):
        with gr.Row():
            img_mt = gr.Image(type="filepath", label="Imagem de Entrada")
            with gr.Column():
                turns_mt = gr.Textbox(lines=5, label="Instruções de Edição (uma por linha)", placeholder="Ex: adicione um chapéu azul\nagora, mude o fundo para uma praia")
                with gr.Accordion("Configurações Avançadas e de Desempenho", open=True):
                    with gr.Row():
                        num_gpus_mt = gr.Slider(label="Número de GPUs", minimum=1, maximum=8, step=1, value=8, info="Use >1 para acelerar a geração com torchrun.")
                        batch_size_mt = gr.Number(label="Batch Size por GPU", value=1, precision=0, info="Para Multi-GPU, o lote total será (GPUs x Batch Size).")
                    resolution_mt = gr.Slider(label="Resolução", minimum=256, maximum=1024, step=128, value=512, info="Maior resolução exige mais VRAM e tempo.")
                    use_vae_slicing_mt = gr.Checkbox(label="Usar VAE Slicing (Economiza VRAM)", value=True)
                    neg_prompt_mt = gr.Textbox(lines=3, label="Prompt Negativo", value=DEFAULT_NEGATIVE_PROMPT)
                    seed_mt = gr.Number(label="Seed (Semente)", value=1, precision=0)
                    steps_mt = gr.Slider(label="Passos de Inferência", minimum=10, maximum=100, step=1, value=50, info="Menos passos = mais rápido.")
                    cfg_mt = gr.Slider(label="Escala de Orientação (CFG)", minimum=1.0, maximum=20.0, step=0.5, value=7.5, info="Quão forte o modelo segue o prompt.")
        run_mt = gr.Button("Executar Edição Multi-Turno", variant="primary")
        gallery_mt = gr.Gallery(label="Imagens Geradas", columns=4, height="auto")
        video_mt = gr.Video(label="Vídeo da Sequência (se disponível)")
        status_mt = gr.Textbox(label="Status da Saída", interactive=False)
        run_mt.click(ui_multi_turn,
                     inputs=[img_mt, turns_mt, neg_prompt_mt, seed_mt, steps_mt, cfg_mt, resolution_mt, use_vae_slicing_mt, num_gpus_mt, batch_size_mt],
                     outputs=[gallery_mt, video_mt, status_mt])

    with gr.Tab("Texto-para-Vídeo"):
        with gr.Row():
            img_vid = gr.Image(type="filepath", label="Frame Inicial")
            with gr.Column():
                prompt_vid = gr.Textbox(lines=2, label="Prompt do Vídeo", placeholder="Ex: um gato andando pela sala")
                with gr.Accordion("Configurações Avançadas e de Desempenho", open=True):
                    with gr.Row():
                        num_gpus_vid = gr.Slider(label="Número de GPUs", minimum=1, maximum=8, step=1, value=8, info="Use >1 para acelerar a geração com torchrun.")
                        batch_size_vid = gr.Number(label="Batch Size por GPU", value=1, precision=0, info="Para Multi-GPU, o lote total será (GPUs x Batch Size).")
                    resolution_vid = gr.Slider(label="Resolução", minimum=256, maximum=1024, step=128, value=512)
                    fps_vid = gr.Slider(label="Frames por Segundo (FPS)", minimum=1, maximum=24, step=1, value=2)
                    use_vae_slicing_vid = gr.Checkbox(label="Usar VAE Slicing (Economiza VRAM)", value=True)
                    neg_prompt_vid = gr.Textbox(lines=3, label="Prompt Negativo", value=DEFAULT_NEGATIVE_PROMPT)
                    seed_vid = gr.Number(label="Seed (Semente)", value=1, precision=0)
                    steps_vid = gr.Slider(label="Passos de Inferência", minimum=10, maximum=100, step=1, value=50)
                    cfg_vid = gr.Slider(label="Escala de Orientação (CFG)", minimum=1.0, maximum=20.0, step=0.5, value=7.5)
        run_vid = gr.Button("Gerar Vídeo", variant="primary")
        video_vid = gr.Video(label="Vídeo Gerado")
        status_vid = gr.Textbox(label="Status da Saída", interactive=False)
        run_vid.click(ui_text_to_video,
                      inputs=[img_vid, prompt_vid, neg_prompt_vid, seed_vid, steps_vid, cfg_vid, resolution_vid, fps_vid, use_vae_slicing_vid, num_gpus_vid, batch_size_vid],
                      outputs=[video_vid, status_vid])

    with gr.Tab("Composição Multi-Conceito"):
        gr.Markdown("Nota: A composição multi-conceito está atualmente configurada para rodar em uma única GPU para garantir estabilidade.")
        with gr.Row():
            with gr.Column(scale=1):
                files_mc = gr.File(file_count="multiple", type="filepath", label="1. Imagens de Conceito")
            with gr.Column(scale=2):
                descs_mc = gr.Textbox(lines=5, label="2. Descrições (uma por linha, na mesma ordem)", placeholder="Ex: <IMG1>: uma foto de um pai\n<IMG2>: uma foto de uma mãe...")
                final_prompt_mc = gr.Textbox(lines=3, label="3. Prompt Final de Composição", placeholder="Ex: Baseado em <IMG0>, <IMG1>..., uma família sorrindo em um retrato...")
        run_mc = gr.Button("Executar Composição", variant="primary")
        gallery_mc = gr.Gallery(label="Imagens Geradas", columns=4, height="auto")
        video_mc = gr.Video(label="Vídeo da Sequência (se disponível)")
        status_mc = gr.Textbox(label="Status da Saída", interactive=False)
        run_mc.click(ui_multi_concept,
                     inputs=[files_mc, descs_mc, final_prompt_mc],
                     outputs=[gallery_mc, video_mc, status_mc])

    # Gatilho de configuração automática no carregamento
    demo.load(fn=setup_auto, outputs=setup_out)

if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=int(os.getenv("PORT", "7860")),
    )