video_service.py

# video_service.py

# --- 1. IMPORTAÇÕES ---
import torch
import numpy as np
import random
import os
import yaml
from pathlib import Path
import imageio
import tempfile
from huggingface_hub import hf_hub_download
import sys
import subprocess

# --- 2. GERENCIAMENTO DE DEPENDÊNCIAS E SETUP ---

def run_setup():
    """Executa o script setup.py para clonar as dependências necessárias."""
    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():
    """Adiciona o repositório clonado ao sys.path para que suas bibliotecas possam ser importadas."""
    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()

# --- 3. 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, LTXMultiScalePipeline
from ltx_video.utils.skip_layer_strategy import SkipLayerStrategy

# --- 4. FUNÇÕES HELPER DE LOG ---
def log_tensor_info(tensor, name="Tensor"):
    if not isinstance(tensor, torch.Tensor):
        print(f"\n[INFO] O item '{name}' não é um tensor para logar.")
        return
    print(f"\n--- Informações do Tensor: {name} ---")
    print(f"  - Shape: {tensor.shape}")
    print(f"  - Dtype: {tensor.dtype}")
    print(f"  - Device: {tensor.device}")
    if tensor.numel() > 0:
        print(f"  - Min valor: {tensor.min().item():.4f}")
        print(f"  - Max valor: {tensor.max().item():.4f}")
        print(f"  - Média: {tensor.mean().item():.4f}")
    else:
        print("  - O tensor está vazio, sem estatísticas.")
    print("------------------------------------------\n")

# --- 5. CLASSE PRINCIPAL DO SERVIÇO ---
class VideoService:
    def __init__(self):
        print("Inicializando VideoService...")
        self.config = self._load_config()
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.last_memory_reserved_mb = 0
        self.pipeline, self.latent_upsampler = self._load_models()
        print(f"Movendo modelos para o dispositivo de inferência: {self.device}")
        self.pipeline.to(self.device)
        if self.latent_upsampler:
            self.latent_upsampler.to(self.device)
        if self.device == "cuda":
            torch.cuda.empty_cache()
            self._log_gpu_memory("Após carregar modelos")
        print("VideoService pronto para uso.")

    def _log_gpu_memory(self, stage_name: str):
        if self.device != "cuda": return
        current_reserved_b = torch.cuda.memory_reserved()
        current_reserved_mb = current_reserved_b / (1024 ** 2)
        total_memory_b = torch.cuda.get_device_properties(0).total_memory
        total_memory_mb = total_memory_b / (1024 ** 2)
        peak_reserved_mb = torch.cuda.max_memory_reserved() / (1024 ** 2)
        delta_mb = current_reserved_mb - self.last_memory_reserved_mb
        print(f"\n--- [LOG DE MEMÓRIA GPU] - {stage_name} ---")
        print(f"  - Uso Atual (Reservado): {current_reserved_mb:.2f} MB / {total_memory_mb:.2f} MB")
        print(f"  - Variação desde o último log: {delta_mb:+.2f} MB")
        if peak_reserved_mb > self.last_memory_reserved_mb:
            print(f"  - Pico de Uso (nesta operação): {peak_reserved_mb:.2f} MB")
        print("--------------------------------------------------\n")
        self.last_memory_reserved_mb = current_reserved_mb

    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 _load_models(self):
        models_dir = "downloaded_models_gradio"
        Path(models_dir).mkdir(parents=True, exist_ok=True)
        LTX_REPO = "Lightricks/LTX-Video"
        distilled_model_path = hf_hub_download(repo_id=LTX_REPO, filename=self.config["checkpoint_path"], local_dir=models_dir, cache_dir=os.getenv("HF_HOME_CACHE"), token=os.getenv("HF_TOKEN"))
        self.config["checkpoint_path"] = distilled_model_path
        spatial_upscaler_path = hf_hub_download(repo_id=LTX_REPO, filename=self.config["spatial_upscaler_model_path"], local_dir=models_dir, cache_dir=os.getenv("HF_HOME_CACHE"), token=os.getenv("HF_TOKEN"))
        self.config["spatial_upscaler_model_path"] = spatial_upscaler_path
        pipeline = create_ltx_video_pipeline(ckpt_path=self.config["checkpoint_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 = None
        if self.config.get("spatial_upscaler_model_path"):
            latent_upsampler = create_latent_upsampler(self.config["spatial_upscaler_model_path"], device="cpu")
        return pipeline, latent_upsampler
        
    def _prepare_conditioning_tensor(self, filepath, height, width, padding_values):
        tensor = load_image_to_tensor_with_resize_and_crop(filepath, height, width)
        tensor = torch.nn.functional.pad(tensor, padding_values)
        return tensor.to(self.device)

    def generate(self, prompt, negative_prompt, mode="text-to-video",
                 start_image_filepath=None,
                 middle_image_filepath=None, middle_frame_number=None, middle_image_weight=1.0,
                 end_image_filepath=None, end_image_weight=1.0,
                 input_video_filepath=None, height=512, width=704, duration=2.0,
                 frames_to_use=9, seed=42, randomize_seed=True, guidance_scale=3.0,
                 improve_texture=True, progress_callback=None):
        if self.device == "cuda":
            torch.cuda.empty_cache()
            torch.cuda.reset_peak_memory_stats()
        self._log_gpu_memory("Início da Geração")

        if mode == "image-to-video" and not start_image_filepath:
            raise ValueError("A imagem de início é obrigatória para o modo image-to-video")
        if mode == "video-to-video" and not input_video_filepath:
            raise ValueError("O vídeo de entrada é 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 = 24.0
        MAX_NUM_FRAMES = 257
        target_frames_rounded = round(duration * FPS)
        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)))
        
        height_padded = ((height - 1) // 32 + 1) * 32
        width_padded = ((width - 1) // 32 + 1) * 32
        padding_values = calculate_padding(height, width, height_padded, width_padded)
        
        generator = torch.Generator(device=self.device).manual_seed(used_seed)
        
        conditioning_items = []
        if mode == "image-to-video":
            start_tensor = self._prepare_conditioning_tensor(start_image_filepath, height, width, padding_values)
            conditioning_items.append(ConditioningItem(start_tensor, 0, 1.0))
            if middle_image_filepath and middle_frame_number is not None:
                middle_tensor = self._prepare_conditioning_tensor(middle_image_filepath, height, width, padding_values)
                safe_middle_frame = max(0, min(int(middle_frame_number), actual_num_frames - 1))
                conditioning_items.append(ConditioningItem(middle_tensor, safe_middle_frame, float(middle_image_weight)))
            if end_image_filepath:
                end_tensor = self._prepare_conditioning_tensor(end_image_filepath, height, width, padding_values)
                last_frame_index = actual_num_frames - 1
                conditioning_items.append(ConditioningItem(end_tensor, last_frame_index, float(end_image_weight)))

        call_kwargs = {
            "prompt": prompt, "negative_prompt": negative_prompt, "height": height_padded, "width": width_padded,
            "num_frames": actual_num_frames, "frame_rate": int(FPS), "generator": generator, "output_type": "pt",
            "conditioning_items": conditioning_items if conditioning_items else 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,
            "is_video": True, "vae_per_channel_normalize": True,
            "mixed_precision": (self.config["precision"] == "mixed_precision"),
            "offload_to_cpu": False, "enhance_prompt": False,
            "skip_layer_strategy": SkipLayerStrategy.AttentionValues
        }

        if mode == "video-to-video":
            call_kwargs["media_items"] = load_media_file(media_path=input_video_filepath, height=height, width=width, max_frames=int(frames_to_use), padding=padding_values).to(self.device)

        result_tensor = None
        if improve_texture:
            if not self.latent_upsampler:
                raise ValueError("Upscaler espacial não carregado.")
            multi_scale_pipeline = LTXMultiScalePipeline(self.pipeline, self.latent_upsampler)
            first_pass_args = self.config.get("first_pass", {}).copy()
            first_pass_args["guidance_scale"] = float(guidance_scale)
            second_pass_args = self.config.get("second_pass", {}).copy()
            second_pass_args["guidance_scale"] = float(guidance_scale)
            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})
            result_tensor = multi_scale_pipeline(**multi_scale_call_kwargs).images
            log_tensor_info(result_tensor, "Resultado da Etapa 2 (Saída do Pipeline Multi-Scale)")
        else:
            single_pass_kwargs = call_kwargs.copy()
            first_pass_config = self.config.get("first_pass", {})
            single_pass_kwargs.update({
                "guidance_scale": float(guidance_scale),
                "stg_scale": first_pass_config.get("stg_scale"),
                "rescaling_scale": first_pass_config.get("rescaling_scale"),
                "skip_block_list": first_pass_config.get("skip_block_list"),
            })

            # --- <INÍCIO DA CORREÇÃO> ---
            if mode == "video-to-video":
                single_pass_kwargs["timesteps"] = [0.7] # CORRIGIDO: Passar como uma lista
                print("[INFO] Modo video-to-video (etapa única): definindo timesteps (força) para [0.7]")
            else:
                single_pass_kwargs["timesteps"] = first_pass_config.get("timesteps")
            # --- <FIM DA CORREÇÃO> ---
            
            print("\n[INFO] Executando pipeline de etapa única...")
            result_tensor = self.pipeline(**single_pass_kwargs).images
        
        pad_left, pad_right, pad_top, pad_bottom = padding_values
        slice_h_end = -pad_bottom if pad_bottom > 0 else None
        slice_w_end = -pad_right if pad_right > 0 else None
        
        result_tensor = result_tensor[:, :, :actual_num_frames, pad_top:slice_h_end, pad_left:slice_w_end]
        log_tensor_info(result_tensor, "Tensor Final (Após Pós-processamento, Antes de Salvar)")

        video_np = (result_tensor[0].permute(1, 2, 3, 0).cpu().float().numpy() * 255).astype(np.uint8)
        temp_dir = tempfile.mkdtemp()
        output_video_path = os.path.join(temp_dir, f"output_{used_seed}.mp4")

        with imageio.get_writer(output_video_path, fps=call_kwargs["frame_rate"], codec='libx264', quality=8) as writer:
            total_frames = len(video_np)
            for i, frame in enumerate(video_np):
                writer.append_data(frame)
                if progress_callback:
                    progress_callback(i + 1, total_frames)
        
        self._log_gpu_memory("Fim da Geração")
        return output_video_path, used_seed

print("Criando instância do VideoService. O carregamento do modelo começará agora...")
video_generation_service = VideoService()