video_service.py

# video_service.py

import torch
import numpy as np
import random
import os
import yaml
from pathlib import Path
import imageio
import tempfile
import sys
import subprocess
import threading
import time
from huggingface_hub import hf_hub_download

# --- LÓGICA DE SETUP E DEPENDÊNCIAS ---

def run_setup():
    setup_script_path = "setup.py"
    if not os.path.exists(setup_script_path):
        print("AVISO: script 'setup.py' não encontrado. Pulando a clonagem de dependências.")
        return
    try:
        print("--- Executando setup.py para garantir que as dependências estão presentes ---")
        subprocess.run([sys.executable, setup_script_path], check=True)
        print("--- Setup concluído com sucesso ---")
    except subprocess.CalledProcessError as e:
        print(f"ERRO CRÍTICO DURANTE O SETUP: 'setup.py' falhou com código {e.returncode}.")
        sys.exit(1)

DEPS_DIR = Path("/data")
LTX_VIDEO_REPO_DIR = DEPS_DIR / "LTX-Video"
if not LTX_VIDEO_REPO_DIR.exists():
    run_setup()

def add_deps_to_path():
    if not LTX_VIDEO_REPO_DIR.exists():
        raise FileNotFoundError(f"Repositório LTX-Video não encontrado em '{LTX_VIDEO_REPO_DIR}'. Execute o setup.")
    if str(LTX_VIDEO_REPO_DIR.resolve()) not in sys.path:
        sys.path.insert(0, str(LTX_VIDEO_REPO_DIR.resolve()))

#add_deps_to_path()

# Importações específicas do modelo
from inference import (
    create_ltx_video_pipeline, create_latent_upsampler,
    load_image_to_tensor_with_resize_and_crop, seed_everething,
    calculate_padding, load_media_file
)
from ltx_video.pipelines.pipeline_ltx_video import ConditioningItem
from ltx_video.utils.skip_layer_strategy import SkipLayerStrategy

# --- CONFIGURAÇÃO DA DISTRIBUIÇÃO DE GPUS ---
GPU_MAPPING = [
    {'base': 'cuda:0', 'upscaler': 'cuda:2'},
    {'base': 'cuda:1', 'upscaler': 'cuda:3'}
]

class VideoService:
    def __init__(self):
        print("Inicializando VideoService (modo Lazy Loading)...")
        self.models_loaded = False
        self.workers = None
        self.config = self._load_config()
        self.models_dir = "downloaded_models"
        self.loading_lock = threading.Lock() # Para evitar que múltiplos usuários iniciem o carregamento ao mesmo tempo

    def _ensure_models_are_loaded(self):
        """Verifica se os modelos estão carregados e os carrega se não estiverem."""
        with self.loading_lock:
            if not self.models_loaded:
                print("Primeira requisição recebida. Iniciando carregamento dos modelos...")
                if torch.cuda.is_available() and torch.cuda.device_count() < 4:
                    raise RuntimeError(f"Este serviço está configurado para 4 GPUs, mas apenas {torch.cuda.device_count()} foram encontradas.")
                
                self._download_model_files()
                self.workers = self._initialize_workers()
                self.models_loaded = True
                print(f"Modelos carregados com sucesso. {len(self.workers)} workers prontos.")

    def _load_config(self):
        config_file_path = LTX_VIDEO_REPO_DIR / "configs" / "ltxv-13b-0.9.8-distilled.yaml"
        with open(config_file_path, "r") as file:
            return yaml.safe_load(file)

    def _download_model_files(self):
        Path(self.models_dir).mkdir(parents=True, exist_ok=True)
        LTX_REPO = "Lightricks/LTX-Video"
        print("Baixando arquivos de modelo (se necessário)...")
        self.distilled_model_path = hf_hub_download(repo_id=LTX_REPO, filename=self.config["checkpoint_path"], local_dir=self.models_dir)
        self.spatial_upscaler_path = hf_hub_download(repo_id=LTX_REPO, filename=self.config["spatial_upscaler_model_path"], local_dir=self.models_dir)
        print("Download de modelos concluído.")

    def _load_models_for_worker(self, base_device, upscaler_device):
        print(f"Carregando modelo base para {base_device} e upscaler para {upscaler_device}")
        pipeline = create_ltx_video_pipeline(
            ckpt_path=self.distilled_model_path, precision=self.config["precision"],
            text_encoder_model_name_or_path=self.config["text_encoder_model_name_or_path"],
            sampler=self.config["sampler"], device="cpu", enhance_prompt=False,
            prompt_enhancer_image_caption_model_name_or_path=self.config["prompt_enhancer_image_caption_model_name_or_path"],
            prompt_enhancer_llm_model_name_or_path=self.config["prompt_enhancer_llm_model_name_or_path"],
        )
        latent_upsampler = create_latent_upsampler(self.spatial_upscaler_path, device="cpu")
        pipeline.to(base_device)
        latent_upsampler.to(upscaler_device)
        return pipeline, latent_upsampler

    def _initialize_workers(self):
        workers = []
        for i, mapping in enumerate(GPU_MAPPING):
            print(f"--- Inicializando Worker {i} ---")
            pipeline, latent_upsampler = self._load_models_for_worker(mapping['base'], mapping['upscaler'])
            workers.append({"id": i, "base_pipeline": pipeline, "latent_upsampler": latent_upsampler, "devices": mapping, "lock": threading.Lock()})
        return workers

    def _acquire_worker(self):
        while True:
            for worker in self.workers:
                if worker["lock"].acquire(blocking=False):
                    print(f"Worker {worker['id']} adquirido para uma nova tarefa.")
                    return worker
            time.sleep(0.1)

    
    def generate(self, prompt, negative_prompt, input_image_filepath=None, input_video_filepath=None, height=512, width=704, mode="text-to-video", duration=2.0, frames_to_use=9, seed=42, randomize_seed=True, guidance_scale=1.0,  # Agora usado corretamente
                 improve_texture=True, progress_callback=None):
        
        self._ensure_models_are_loaded()
        worker = self._acquire_worker()
        base_device = worker['devices']['base']
        upscaler_device = worker['devices']['upscaler']
        
        try:
            # Validações alinhadas com app-20.py
            if mode == "image-to-video" and not input_image_filepath:
                raise ValueError("Caminho da imagem obrigatório para o modo image-to-video")
            if mode == "video-to-video" and not input_video_filepath:
                raise ValueError("Caminho do vídeo obrigatório para o modo video-to-video")
            
            used_seed = random.randint(0, 2**32 - 1) if randomize_seed else int(seed)
            seed_everething(used_seed)
            
            FPS = 30.0  # Alinhado com app-20.py
            MAX_NUM_FRAMES = 257
            target_frames_ideal = duration * FPS
            target_frames_rounded = round(target_frames_ideal)
            if target_frames_rounded < 1: target_frames_rounded = 1
            n_val = round(float(target_frames_rounded - 1.0) / 8.0)
            actual_num_frames = max(9, min(MAX_NUM_FRAMES, int(n_val * 8 + 1)))
            
            actual_height = int(height)
            actual_width = int(width)
            height_padded = (actual_height - 1) // 32 * 32 + 32
            width_padded = (actual_width - 1) // 32 * 32 + 32
            num_frames_padded = (actual_num_frames - 2) // 8 * 8 + 1  # Alinhamento exato com app-20.py
            if num_frames_padded != actual_num_frames:
                print(f"Warning: actual_num_frames {actual_num_frames} and num_frames_padded {num_frames_padded} differ. Using num_frames_padded for pipeline.")
            
            padding_values = calculate_padding(actual_height, actual_width, height_padded, width_padded)
            pad_left, pad_right, pad_top, pad_bottom = padding_values
            
            # Kwargs base alinhados
            call_kwargs = {
                "prompt": prompt,
                "negative_prompt": negative_prompt,
                "height": height_padded,
                "width": width_padded,
                "num_frames": num_frames_padded,
                "framerate": int(FPS),
                "generator": torch.Generator(device=base_device).manual_seed(used_seed),
                "output_type": "pt",
                "conditioning_items": None,
                "media_items": None,
                "decode_timestep": self.config['decode_timestep'],
                "decode_noise_scale": self.config['decode_noise_scale'],
                "stochastic_sampling": self.config['stochastic_sampling'],
                "image_cond_noise_scale": 0.15,  # Alinhado
                "is_video": True,
                "vae_per_channel_normalize": True,
                "mixed_precision": self.config['precision'] + " mixed_precision",
                "offload_to_cpu": False,
                "enhance_prompt": False,
            }
            
            # Estratégia de skip layer alinhada
            stg_mode_str = self.config.get('stg_mode', 'attention_values')
            if stg_mode_str.lower() in ['stgav', 'attentionvalues']:
                call_kwargs['skip_layer_strategy'] = SkipLayerStrategy.AttentionValues
            # ... (adicionar outros elif como no app-20.py)
            
            # Conditioning para modos
            if mode == "image-to-video" and input_image_filepath:
                media_tensor = load_image_to_tensor_with_resize_and_crop(input_image_filepath, actual_height, actual_width)
                media_tensor = torch.nn.functional.pad(media_tensor, padding_values)
                call_kwargs['conditioning_items'] = ConditioningItem(media_tensor.to(base_device), 0, 1.0)
            elif mode == "video-to-video" and input_video_filepath:
                call_kwargs['media_items'] = load_media_file(media_path=input_video_filepath, height=actual_height, width=actual_width, max_frames=int(frames_to_use), padding=padding_values).to(base_device)
            
            result_images_tensor = None
            if improve_texture:
                # Alinhamento exato: Use LTXMultiScalePipeline como no app-20.py
                active_latent_upsampler = worker['latent_upsampler']
                if not active_latent_upsampler:
                    raise ValueError("Spatial upscaler model not loaded or improve_texture not selected, cannot use multi-scale.")
                
                multi_scale_pipeline_obj = LTXMultiScalePipeline(worker['base_pipeline'], active_latent_upsampler)
                
                first_pass_args = self.config.get('first_pass', {}).copy()
                first_pass_args['guidance_scale'] = float(guidance_scale)  # Override UI
                first_pass_args.pop('num_inference_steps', None)
                
                second_pass_args = self.config.get('second_pass', {}).copy()
                second_pass_args['guidance_scale'] = float(guidance_scale)  # Override UI
                second_pass_args.pop('num_inference_steps', None)
                
                multi_scale_call_kwargs = call_kwargs.copy()
                multi_scale_call_kwargs.update({
                    "downscale_factor": self.config['downscale_factor'],
                    "first_pass": first_pass_args,
                    "second_pass": second_pass_args,
                })
                
                print(f"Calling multi-scale pipeline eff. HxW {actual_height}x{actual_width}, Frames {actual_num_frames} - Padded {num_frames_padded} on {base_device}")
                result_images_tensor = multi_scale_pipeline_obj(**multi_scale_call_kwargs).images
            else:
                # Single-pass alinhado
                single_pass_call_kwargs = call_kwargs.copy()
                first_pass_config_from_yaml = self.config.get('first_pass', {})
                single_pass_call_kwargs['timesteps'] = first_pass_config_from_yaml.get('timesteps')
                single_pass_call_kwargs['guidance_scale'] = float(guidance_scale)  # Override UI
                single_pass_call_kwargs['stg_scale'] = first_pass_config_from_yaml.get('stg_scale')
                single_pass_call_kwargs['rescaling_scale'] = first_pass_config_from_yaml.get('rescaling_scale')
                single_pass_call_kwargs['skip_block_list'] = first_pass_config_from_yaml.get('skip_block_list')
                single_pass_call_kwargs.pop('num_inference_steps', None)
                single_pass_call_kwargs.pop('first_pass', None)
                single_pass_call_kwargs.pop('second_pass', None)
                single_pass_call_kwargs.pop('downscale_factor', None)
                
                print(f"Calling base pipeline padded HxW {height_padded}x{width_padded}, Frames {actual_num_frames} - Padded {num_frames_padded} on {base_device}")
                result_images_tensor = worker['base_pipeline'](**single_pass_call_kwargs).images
            
            if result_images_tensor is None:
                raise ValueError("Generation failed.")
            
            # Slicing e salvamento alinhados
            slice_h_end = -pad_bottom if pad_bottom > 0 else None
            slice_w_end = -pad_right if pad_right > 0 else None
            result_images_tensor = result_images_tensor[:, :, :actual_num_frames, pad_top:slice_h_end, pad_left:slice_w_end]
            video_np = result_images_tensor[0].permute(1, 2, 3, 0).cpu().float().numpy()
            video_np = np.clip(video_np, 0, 1) * 255.0
            video_np = video_np.astype(np.uint8)
            
            temp_dir = tempfile.mkdtemp()
            output_video_path = os.path.join(temp_dir, f"output_{used_seed}.mp4")
            try:
                with imageio.get_writer(output_video_path, fps=call_kwargs['framerate'], macro_block_size=1) as video_writer:
                    for frame_idx in range(video_np.shape[0]):
                        if progress_callback:
                            progress_callback(frame_idx / video_np.shape[0], desc="Saving video")
                        video_writer.append_data(video_np[frame_idx])
            except Exception as e:
                print(f"Error saving video with macro_block_size=1: {e}")
                with imageio.get_writer(output_video_path, fps=call_kwargs['framerate'], format='FFMPEG', codec='libx264', quality=8) as video_writer:
                    for frame_idx in range(video_np.shape[0]):
                        if progress_callback:
                            progress_callback(frame_idx / video_np.shape[0], desc="Saving video fallback ffmpeg")
                        video_writer.append_data(video_np[frame_idx])
            
            return output_video_path, used_seed
        
        except Exception as e:
            print(f"!!!!!!!! ERRO no Worker {worker['id']}: {e} !!!!!!!!")
            raise e
        finally:
            print(f"Worker {worker['id']} Tarefa finalizada. Limpando cache e liberando worker...")
            with torch.cuda.device(base_device):
                torch.cuda.empty_cache()
            with torch.cuda.device(upscaler_device):
                torch.cuda.empty_cache()
            worker['lock'].release()

# A instância do serviço é criada aqui, mas os modelos só serão carregados no primeiro clique.
video_generation_service = VideoService()