app.py

# app.py (Versão Unificada LTX + SeedVR)

# --- 0. WARNINGS E SETUP ---
import warnings
warnings.filterwarnings("ignore")
import gradio as gr
from PIL import Image
import os
import imageio
import multiprocessing as mp

# --- 1. IMPORTAÇÃO DOS DOIS SERVIÇOS ---

# Importa o serviço de geração do LTX
try:
    from api.ltx_server import video_generation_service
except ImportError:
    raise ImportError("Falha ao importar o 'video_generation_service' de 'api.ltx_server'.")

# Importa o serviço de upscale do SeedVR
try:
    from api.seedvr_server import SeedVRServer
except ImportError:
    raise ImportError("Falha ao importar o 'SeedVRServer' de 'api.seedvr_server'.")

# --- 2. INICIALIZAÇÃO DOS SERVIÇOS ---
print("🚀 Inicializando o serviço de geração LTX-Video...")
ltx_service = video_generation_service
print("✅ Serviço LTX-Video pronto.")

print("🚀 Inicializando o serviço de upscale SeedVR...")
seedvr_service = SeedVRServer()
print("✅ Serviço SeedVR pronto.")


# --- 3. FUNÇÕES HELPER DA UI (sem grandes alterações) ---
TARGET_FIXED_SIDE = 768
MIN_DIM_SLIDER = 256
MAX_IMAGE_SIZE = 1280

def calculate_new_dimensions(orig_w, orig_h):
    # ... (código inalterado)
    if orig_w == 0 or orig_h == 0: return int(TARGET_FIXED_SIDE), int(TARGET_FIXED_SIDE)
    if orig_w >= orig_h:
        new_h, aspect_ratio = TARGET_FIXED_SIDE, orig_w / orig_h
        new_w = round((new_h * aspect_ratio) / 32) * 32
    else:
        new_w, aspect_ratio = TARGET_FIXED_SIDE, orig_h / orig_w
        new_h = round((new_w * aspect_ratio) / 32) * 32
    return int(max(MIN_DIM_SLIDER, min(new_h, MAX_IMAGE_SIZE))), int(max(MIN_DIM_SLIDER, min(new_w, MAX_IMAGE_SIZE)))

def handle_media_upload_for_dims(filepath, current_h, current_w):
    # ... (código inalterado)
    if not filepath or not os.path.exists(str(filepath)): return gr.update(value=current_h), gr.update(value=current_w)
    try:
        if str(filepath).lower().endswith(('.png', '.jpg', '.jpeg', '.webp')):
            with Image.open(filepath) as img: orig_w, orig_h = img.size
        else:
            with imageio.get_reader(filepath) as reader: meta = reader.get_meta_data(); orig_w, orig_h = meta.get('size', (current_w, current_h))
        new_h, new_w = calculate_new_dimensions(orig_w, orig_h)
        return gr.update(value=new_h), gr.update(value=new_w)
    except Exception as e:
        print(f"Erro ao processar mídia: {e}")
        return gr.update(value=current_h), gr.update(value=current_w)

def update_frame_slider(duration):
    # ... (código inalterado)
    max_frames = int(duration * 24.0)
    new_value = 48 if max_frames >= 48 else max_frames // 2
    return gr.update(maximum=max_frames, value=new_value)

# --- 4. FUNÇÕES DE CALLBACK PARA CADA ETAPA ---

# ETAPA 1: Geração com LTX
def run_ltx_generation(
    prompt, negative_prompt, mode, start_image, middle_image, middle_frame,
    middle_weight, end_image, end_weight, input_video, height, width,
    duration, frames_to_use, seed, randomize_seed, guidance_scale, improve_texture,
    progress=gr.Progress(track_tqdm=True)
):
    # Limpa os resultados da etapa anterior (SeedVR) e mostra o container de upscale
    yield (
        gr.update(interactive=False, value="Generating with LTX..."), # Botão LTX
        None, # Vídeo LTX
        gr.update(visible=False), # Container de upscale
        gr.update(value=None, visible=False), # Log SeedVR
        gr.update(value=None, visible=False)  # Vídeo SeedVR
    )
    
    try:
        def progress_handler(step, total_steps):
            progress(step / total_steps, desc="Saving LTX video...")

        output_path, used_seed = ltx_service.generate(
            prompt=prompt, negative_prompt=negative_prompt, mode=mode,
            start_image_filepath=start_image, middle_image_filepath=middle_image,
            middle_frame_number=middle_frame, middle_image_weight=middle_weight,
            end_image_filepath=end_image, end_image_weight=end_weight,
            input_video_filepath=input_video, height=int(height), width=int(width),
            duration=float(duration), frames_to_use=int(frames_to_use), seed=int(seed),
            randomize_seed=bool(randomize_seed), guidance_scale=float(guidance_scale),
            improve_texture=bool(improve_texture), progress_callback=progress_handler
        )
        
        # Mostra o resultado do LTX e o container de upscale
        yield (
            gr.update(interactive=True, value="Generate Video"),
            gr.update(value=output_path),
            gr.update(visible=True), # Mostra o container de upscale
            gr.update(visible=False),
            gr.update(value=None, visible=False)
        )
    except Exception as e:
        gr.Error(f"LTX Generation Failed: {e}")
        yield gr.update(interactive=True, value="Generate Video"), None, gr.update(visible=False), None, None

# ETAPA 2: Upscale com SeedVR
def run_seedvr_upscale(
    ltx_video_path: str, # Recebe o caminho do vídeo gerado
    resolution: str,
    sp_size: int,
    progress=gr.Progress(track_tqdm=True)
):
    # Desabilita o botão de upscale e mostra os componentes de log/vídeo
    yield (
        gr.update(interactive=False, value="Upscaling with SeedVR..."),
        gr.update(value="▶ Starting SeedVR upscale process...\n", visible=True),
        gr.update(value=None, visible=False)
    )
    
    if not ltx_video_path:
        gr.Warning("No video from LTX to process.")
        yield gr.update(interactive=True), gr.update(value="No input video.", visible=True), None
        return

    log_buffer = ["▶ Starting SeedVR upscale process...\n"]
    last_log_message = ""

    try:
        def progress_callback_wrapper(step: float, desc: str):
            nonlocal last_log_message
            if desc != last_log_message:
                log_buffer.append(f"{desc}\n")
                last_log_message = desc
            progress(step, desc=desc)

        final_video_path = seedvr_service.run_inference_direct(
            file_path=ltx_video_path,
            seed=42,
            res_h=int(resolution),
            res_w=int(resolution),
            sp_size=int(sp_size),
            fps=None, # Usa o FPS original
            progress=progress_callback_wrapper
        )
        
        log_buffer.append("✅ SeedVR Upscale complete!")
        
        yield (
            gr.update(interactive=True, value="Refine and Upscale with SeedVR"),
            ''.join(log_buffer),
            gr.update(value=final_video_path, visible=True)
        )
        
    except Exception as e:
        error_message = f"❌ SeedVR Upscale Failed: {e}"
        gr.Error(error_message)
        log_buffer.append(f"\n{error_message}")
        yield gr.update(interactive=True), ''.join(log_buffer), None

# --- 5. DEFINIÇÃO DA INTERFACE GRADIO UNIFICADA ---
css = "#col-container { margin: 0 auto; max-width: 900px; }"
with gr.Blocks(css=css, theme=gr.themes.Soft(primary_hue="blue")) as demo:
    gr.Markdown("# LTX-Video + SeedVR: Pipeline de Geração e Refinamento")
    
    # --- ETAPA 1: GERAÇÃO (LTX) ---
    gr.Markdown("## Etapa 1: Gerar Vídeo Base com LTX-Video")
    with gr.Row():
        with gr.Column():
            # ... (Toda a UI do LTX-Video que você já tem vai aqui)
            with gr.Tab("image-to-video (Keyframes)") as image_tab:
                # ...
                i2v_prompt = gr.Textbox(label="Prompt", value="Uma bela transição entre as imagens", lines=2)
                start_image_i2v = gr.Image(label="Imagem de Início", type="filepath")
                # ... etc ...
                i2v_button = gr.Button("Generate Image-to-Video", variant="primary")

            with gr.Tab("text-to-video") as text_tab:
                t2v_prompt = gr.Textbox(label="Prompt", value="A majestic dragon flying over a medieval castle", lines=3)
                t2v_button = gr.Button("Generate Text-to-Video", variant="primary")
                
            # ... (outras abas e configurações do LTX)

            duration_input = gr.Slider(label="Video Duration (seconds)", minimum=1, maximum=30, value=8, step=0.5)
            improve_texture = gr.Checkbox(label="Improve Texture (multi-scale)", value=True)
            
        with gr.Column():
            ltx_output_video = gr.Video(label="LTX Generated Video", interactive=False)

    # --- ETAPA 2: UPSCALE (SeedVR) ---
    with gr.Box(visible=False) as upscale_container:
        gr.Markdown("---")
        gr.Markdown("## Etapa 2: Refinar e Aumentar Resolução com SeedVR (Multi-GPU)")
        
        with gr.Row():
            with gr.Column(scale=1):
                seedvr_resolution = gr.Dropdown(
                    label="Target Resolution",
                    choices=["720", "960", "1024", "2048"],
                    value="1024"
                )
                seedvr_sp_size = gr.Slider(
                    label="Frames per Batch (sp_size)",
                    minimum=1, maximum=16, step=1, value=4
                )
                upscale_button = gr.Button("Refine and Upscale with SeedVR", variant="primary", icon="💎")

            with gr.Column(scale=2):
                seedvr_log_window = gr.Textbox(label="SeedVR Inference Log 📝", lines=8, visible=False, autoscroll=True)
                seedvr_output_video = gr.Video(label="Final Upscaled Video", visible=False)

    # --- Acorddion de Configs Avançadas (como antes) ---
    with gr.Accordion("Advanced settings", open=False):
        # ... (todos os inputs avançados do LTX)
        mode = gr.Dropdown(["text-to-video", "image-to-video", "video-to-video"], label="task", value="image-to-video", visible=False)
        # ...

    # --- LÓGICA DE EVENTOS ---
    # ... (eventos de upload e seleção de abas do LTX, como antes) ...

    # Evento do botão de geração do LTX
    ltx_inputs = [...] # Sua lista completa de inputs para o LTX
    i2v_button.click(
        fn=run_ltx_generation,
        inputs=ltx_inputs,
        outputs=[i2v_button, ltx_output_video, upscale_container, seedvr_log_window, seedvr_output_video]
    )
    # (faça o mesmo para t2v_button e v2v_button)

    # Evento do botão de upscale do SeedVR
    upscale_button.click(
        fn=run_seedvr_upscale,
        inputs=[ltx_output_video, seedvr_resolution, seedvr_sp_size],
        outputs=[upscale_button, seedvr_log_window, seedvr_output_video]
    )

if __name__ == "__main__":
    mp.set_start_method('spawn', force=True)
    demo.queue().launch(debug=True, share=False)