Update api/ltx_server_refactored.py

api/ltx_server_refactored.py

@@ -33,6 +33,7 @@ import shutil
 import contextlib
 import time
 import traceback
+from api.gpu_manager import gpu_manager
 from einops import rearrange
 import torch.nn.functional as F
 from managers.vae_manager import vae_manager_singleton
@@ -101,6 +102,9 @@ class VideoService:
     def __init__(self):
         t0 = time.perf_counter()
         print("[DEBUG] Inicializando VideoService...")
+        self.device = gpu_manager.get_ltx_device()
+        print(f"[DEBUG] LTX foi alocado para o dispositivo: {self.device}")
+        
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
         self.config = self._load_config()
         self.pipeline, self.latent_upsampler = self._load_models()
@@ -139,7 +143,21 @@ class VideoService:
             self._log_gpu_memory("Após finalize")
         except Exception as e:
             print(f"[DEBUG] Log GPU pós-finalize falhou: {e}")
-            
+
+    def move_to_device(self, device):
+        """Move os modelos do pipeline para o dispositivo especificado."""
+        print(f"[LTX] Movendo modelos para {device}...")
+        self.pipeline.to(device)
+        if self.latent_upsampler:
+            self.latent_upsampler.to(device)
+        self.device = device
+        
+    def move_to_cpu(self):
+        """Move os modelos para a CPU para liberar VRAM."""
+        self.move_to_device(torch.device("cpu"))
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+    
     def _load_models(self):
         t0 = time.perf_counter()
         LTX_REPO = "Lightricks/LTX-Video"
@@ -245,7 +263,7 @@ class VideoService:
             conditioning_items.append(ConditioningItem(tensor, safe_frame, float(weight)))
         return conditioning_items
 
-    def generate_low(self, prompt, negative_prompt, height, width, duration, guidance_scale, seed, conditioning_items=None):
+    def generate_low_old(self, prompt, negative_prompt, height, width, duration, guidance_scale, seed, conditioning_items=None):
         used_seed = random.randint(0, 2**32 - 1) if seed is None else int(seed)
         seed_everething(used_seed)
         FPS = 24.0
@@ -282,7 +300,191 @@ class VideoService:
             torch.cuda.empty_cache()
             torch.cuda.ipc_collect()
             self.finalize(keep_paths=[])
+
+    def _generate_single_chunk_low(self, prompt, negative_prompt, height, width, num_frames, guidance_scale, seed, initial_latent_condition=None, image_conditions=None, ltx_configs_override=None):
+        """
+        [NÓ DE GERAÇÃO]
+        Gera um ÚNICO chunk de latentes brutos. Esta é a unidade de trabalho fundamental.
+        """
+        # (Esta função auxiliar permanece a mesma da nossa última versão, com a lógica de override)
+        print("\n" + "-"*20 + " INÍCIO: _generate_single_chunk_low " + "-"*20)
+        height_padded = ((height - 1) // 8 + 1) * 8
+        width_padded = ((width - 1) // 8 + 1) * 8
+        generator = torch.Generator(device=self.device).manual_seed(seed)
+        
+        downscale_factor = self.config.get("downscale_factor", 0.6666666)
+        vae_scale_factor = self.pipeline.vae_scale_factor
+        
+        x_width = int(width_padded * downscale_factor)
+        downscaled_width = x_width - (x_width % vae_scale_factor)
+        x_height = int(height_padded * downscale_factor)
+        downscaled_height = x_height - (x_height % vae_scale_factor)
+        
+        all_conditions = []
+        if image_conditions: all_conditions.extend(image_conditions)
+        if initial_latent_condition: all_conditions.append(initial_latent_condition)
+
+        first_pass_config = self.config.get("first_pass", {}).copy()
+
+        if ltx_configs_override:
+            print("[DEBUG] Sobrepondo configurações do LTX com valores da UI...")
+            if "first_pass_num_inference_steps" in ltx_configs_override:
+                first_pass_config["num_inference_steps"] = ltx_configs_override["first_pass_num_inference_steps"]
+            if "first_pass_guidance_scale" in ltx_configs_override:
+                max_val = max(first_pass_config.get("guidance_scale", [1]))
+                new_max_val = ltx_configs_override["first_pass_guidance_scale"]
+                first_pass_config["guidance_scale"] = [new_max_val if x==max_val else x for x in first_pass_config["guidance_scale"]]
+        
+        first_pass_kwargs = {
+            "prompt": prompt, "negative_prompt": negative_prompt, "height": downscaled_height, "width": downscaled_width,
+            "num_frames": num_frames, "frame_rate": 24, "generator": generator, "output_type": "latent",
+            "conditioning_items": all_conditions if all_conditions else None,
+            **first_pass_config
+        }
+        # Removido guidance_scale daqui pois agora está dentro do first_pass_config
+        if "guidance_scale" in first_pass_kwargs:
+            del first_pass_kwargs['guidance_scale']
+
+        with torch.autocast(device_type="cuda", dtype=self.runtime_autocast_dtype, enabled=self.device.type == 'cuda'):
+            latents_bruto = self.pipeline(**first_pass_kwargs).images
+            log_tensor_info(latents_bruto, f"Latente Bruto Gerado para: '{prompt[:40]}...'")
+        
+        print("-" * 20 + " FIM: _generate_single_chunk_low " + "-"*20)
+        return latents_bruto
+
+    def generate_narrative_low(self, prompt: str, negative_prompt, height, width, duration, guidance_scale, seed, initial_image_conditions=None, overlap_frames: int = 8, ltx_configs_override: dict = None):
+        """
+        [ORQUESTRADOR NARRATIVO]
+        Gera um vídeo em múltiplos chunks sequenciais a partir de um prompt com várias linhas.
+        """
+        print("\n" + "="*80)
+        print("======           INICIANDO GERAÇÃO NARRATIVA EM CHUNKS (LOW-RES)           ======")
+        print("="*80)
+        
+        used_seed = random.randint(0, 2**32 - 1) if seed is None else int(seed)
+        seed_everething(used_seed)
+        FPS = 24.0
+
+        prompt_list = [p.strip() for p in prompt.splitlines() if p.strip()]
+        num_chunks = len(prompt_list)
+        if num_chunks == 0: raise ValueError("O prompt está vazio ou não contém linhas válidas.")
+        
+        total_actual_frames = max(9, int(round((round(duration * FPS) - 1) / 8.0) * 8 + 1))
+        
+        if num_chunks > 1:
+            total_blocks = (total_actual_frames - 1) // 8
+            blocks_per_chunk = total_blocks // num_chunks
+            blocks_last_chunk = total_blocks - (blocks_per_chunk * (num_chunks - 1))
+            frames_per_chunk = blocks_per_chunk * 8 + 1
+            frames_per_chunk_last = blocks_last_chunk * 8 + 1
+        else:
+            frames_per_chunk = total_actual_frames
+            frames_per_chunk_last = total_actual_frames
+        
+        frames_per_chunk = max(9, frames_per_chunk)
+        frames_per_chunk_last = max(9, frames_per_chunk_last)
+        
+        poda_latents_num = overlap_frames // self.pipeline.video_scale_factor if self.pipeline.video_scale_factor > 0 else 0
+        
+        latentes_chunk_video = []
+        condition_item_latent_overlap = None
+        temp_dir = tempfile.mkdtemp(prefix="ltxv_narrative_"); self._register_tmp_dir(temp_dir)
+        results_dir = "/app/output"; os.makedirs(results_dir, exist_ok=True)
+        
+        for i, chunk_prompt in enumerate(prompt_list):
+            print(f"\n--- Gerando Chunk Narrativo {i+1}/{num_chunks}: '{chunk_prompt}' ---")
+
+            current_image_conditions = []
+            if initial_image_conditions:
+                cond_item_original = initial_image_conditions[0]
+                if i == 0:
+                    current_image_conditions.append(cond_item_original)
+                else:
+                    cond_item_fraco = ConditioningItem(
+                        media_item=cond_item_original.media_item, media_frame_number=0, conditioning_strength=0.1
+                    )
+                    current_image_conditions.append(cond_item_fraco)
+
+            num_frames_para_gerar = frames_per_chunk_last if i == num_chunks - 1 else frames_per_chunk
+            if i > 0 and poda_latents_num > 0:
+                num_frames_para_gerar += overlap_frames
             
+            latentes_bruto = self._generate_single_chunk_low(
+                prompt=chunk_prompt, negative_prompt=negative_prompt, height=height, width=width,
+                num_frames=num_frames_para_gerar, guidance_scale=guidance_scale, seed=used_seed + i,
+                initial_latent_condition=condition_item_latent_overlap, image_conditions=current_image_conditions,
+                ltx_configs_override=ltx_configs_override
+            )
+            
+            if i > 0 and poda_latents_num > 0:
+                 latentes_bruto = latentes_bruto[:, :, poda_latents_num:, :, :]
+
+            latentes_podado = latentes_bruto.clone().detach()
+            if i < num_chunks - 1 and poda_latents_num > 0:
+                latentes_podado = latentes_bruto[:, :, :-poda_latents_num, :, :].clone()
+                overlap_latents = latentes_bruto[:, :, -poda_latents_num:, :, :].clone()
+                condition_item_latent_overlap = ConditioningItem(
+                    media_item=overlap_latents, media_frame_number=0, conditioning_strength=1.0
+                )
+            latentes_chunk_video.append(latentes_podado)
+
+        print("\n--- Finalizando Narrativa: Concatenando chunks ---")
+        final_latents = torch.cat(latentes_chunk_video, dim=2)
+        log_tensor_info(final_latents, "Tensor de Latentes Final Concatenado")
+        
+        tensor_path = os.path.join(results_dir, f"latents_narrative_{used_seed}.pt")
+        torch.save(final_latents.cpu(), tensor_path)
+        
+        with torch.autocast(device_type="cuda", dtype=self.runtime_autocast_dtype, enabled=self.device.type == 'cuda'):
+            pixel_tensor = vae_manager_singleton.decode(final_latents, decode_timestep=float(self.config.get("decode_timestep", 0.05)))
+            video_path = self._save_and_log_video(pixel_tensor, "narrative_video", FPS, temp_dir, results_dir, used_seed)
+        
+        self.finalize(keep_paths=[video_path, tensor_path])
+        return video_path, tensor_path, used_seed
+        
+    def generate_single_low(self, prompt: str, negative_prompt, height, width, duration, guidance_scale, seed, initial_image_conditions=None, ltx_configs_override: dict = None):
+        """
+        [ORQUESTRADOR SIMPLES]
+        Gera um vídeo completo em um único chunk. Ideal para prompts simples e curtos.
+        """
+        print("\n" + "="*80)
+        print("======             INICIANDO GERAÇÃO SIMPLES EM CHUNK ÚNICO (LOW-RES)             ======")
+        print("="*80)
+
+        used_seed = random.randint(0, 2**32 - 1) if seed is None else int(seed)
+        seed_everething(used_seed)
+        FPS = 24.0
+
+        total_actual_frames = max(9, int(round((round(duration * FPS) - 1) / 8.0) * 8 + 1))
+        
+        temp_dir = tempfile.mkdtemp(prefix="ltxv_single_"); self._register_tmp_dir(temp_dir)
+        results_dir = "/app/output"; os.makedirs(results_dir, exist_ok=True)
+
+        # Chama a função de geração de chunk único para fazer todo o trabalho
+        final_latents = self._generate_single_chunk_low(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            height=height, width=width,
+            num_frames=total_actual_frames,
+            guidance_scale=guidance_scale,
+            seed=used_seed,
+            image_conditions=initial_image_conditions,
+            ltx_configs_override=ltx_configs_override
+        )
+
+        print("\n--- Finalizando Geração Simples: Salvando e decodificando ---")
+        log_tensor_info(final_latents, "Tensor de Latentes Final")
+
+        tensor_path = os.path.join(results_dir, f"latents_single_{used_seed}.pt")
+        torch.save(final_latents.cpu(), tensor_path)
+
+        with torch.autocast(device_type="cuda", dtype=self.runtime_autocast_dtype, enabled=self.device.type == 'cuda'):
+            pixel_tensor = vae_manager_singleton.decode(final_latents, decode_timestep=float(self.config.get("decode_timestep", 0.05)))
+            video_path = self._save_and_log_video(pixel_tensor, "single_video", FPS, temp_dir, results_dir, used_seed)
+
+        self.finalize(keep_paths=[video_path, tensor_path])
+        return video_path, tensor_path, used_seed
+        
     def generate_upscale_denoise(self, latents_path, prompt, negative_prompt, guidance_scale, seed):
             used_seed = random.randint(0, 2**32 - 1) if seed is None else int(seed)
             seed_everething(used_seed)
@@ -330,8 +532,6 @@ class VideoService:
             video_path = self._save_and_log_video(pixel_tensor, "refined_video", 24.0, temp_dir, results_dir, used_seed)
             return video_path, tensor_path    
 
-            
-
     def encode_mp4(self, latents_path: str, fps: int = 24):
         latents = torch.load(latents_path)
         seed = random.randint(0, 99999)
@@ -362,6 +562,8 @@ class VideoService:
         
 
 # --- INSTANCIAÇÃO DO SERVIÇO ---
-print("Criando instância do VideoService. O carregamento do modelo começará agora...")
+print("Criando instância do VideoService...")
 video_generation_service = VideoService()
-print("Instância do VideoService pronta para uso.")
+print("Instância do VideoService pronta.")
+self.device = gpu_manager.get_ltx_device()
+        print(f"[DEBUG] LTX foi alocado para o dispositivo: {self.device}")