Update api/ltx_server_refactored.py

api/ltx_server_refactored.py

@@ -444,8 +444,130 @@ class VideoService:
         downscale_factor = self.config.get("downscale_factor", 0.6666666)
         vae_scale_factor = self.pipeline.vae_scale_factor
 
+
         # --- <INÍCIO DA LÓGICA DE CÁLCULO EXATA> ---
         # Replica a fórmula da LTXMultiScalePipeline
         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)
+        print(f"[DEBUG] First Pass Dims: Original Pad ({width_padded}x{height_padded}) -> Downscaled ({downscaled_width}x{downscaled_height})")
+        # --- <FIM DA LÓGICA DE CÁLCULO EXATA> ---
+
+        first_pass_kwargs = {
+            "prompt": prompt, "negative_prompt": negative_prompt, "height": downscaled_height, "width": downscaled_width,
+            "num_frames": actual_num_frames, "frame_rate": int(FPS), "generator": generator, "output_type": "latent",
+            "conditioning_items": conditioning_items, "guidance_scale": float(guidance_scale),
+            **(self.config.get("first_pass", {}))
+        }
+
+        with torch.autocast(device_type="cuda", dtype=self.runtime_autocast_dtype, enabled=self.device == 'cuda'):
+            latents = self.pipeline(**first_pass_kwargs).images
+            log_tensor_info(latents, "Latentes Low-Res Gerados")
+
+            pixel_tensor = vae_manager_singleton.decode(latents.clone(), decode_timestep=float(self.config.get("decode_timestep", 0.05)))
+            video_path = self._save_and_log_video(pixel_tensor, "low_res_video", FPS, temp_dir, results_dir, used_seed)
+            del pixel_tensor
+
+            latents_cpu = latents.detach().to("cpu")
+            tensor_path = os.path.join(results_dir, f"latents_low_res_{used_seed}.pt")
+            torch.save(latents_cpu, tensor_path)
+            print(f"[DEBUG] Tensor latente de baixa resolução salvo em: {tensor_path}")
+        
+        self._log_gpu_memory("Fim da Geração Low-Res")
+        return video_path, tensor_path, used_seed
+
+    def generate_upscale_denoise(self, latents_path, prompt, negative_prompt, guidance_scale, seed):
+        print("\n--- INICIANDO ETAPA 2: UPSCALE E REFINAMENTO ---")
+        self._log_gpu_memory("Início do Upscale/Denoise")
+
+        used_seed = random.randint(0, 2**32 - 1) if seed is None else int(seed)
+        seed_everething(used_seed)
+
+        temp_dir = tempfile.mkdtemp(prefix="ltxv_up_"); self._register_tmp_dir(temp_dir)
+        results_dir = "/app/output"; os.makedirs(results_dir, exist_ok=True)
+        
+        latents_low = torch.load(latents_path).to(self.device)
+        log_tensor_info(latents_low, "Latentes Low-Res Carregados")
+        
+        with torch.autocast(device_type="cuda", dtype=self.runtime_autocast_dtype, enabled=self.device == 'cuda'):
+            upsampled_latents = self._upsample_latents_internal(latents_low)
+            upsampled_latents = adain_filter_latent(latents=upsampled_latents, reference_latents=latents_low)
+            del latents_low; torch.cuda.empty_cache()
+            
+            total_frames = upsampled_latents.shape[2]
+            mid_point = total_frames // 2
+            chunk1 = upsampled_latents[:, :, :mid_point, :, :]
+            chunk2 = upsampled_latents[:, :, mid_point:, :, :]
+            
+            final_latents_list = []
+            
+            for i, chunk in enumerate([chunk1, chunk2]):
+                if chunk.shape[2] == 0: continue
+                print(f"  - Refinando chunk {i+1}/{2} com {chunk.shape[2]} frames")
+                second_pass_height = chunk.shape[3] * self.pipeline.vae_scale_factor
+                second_pass_width = chunk.shape[4] * self.pipeline.vae_scale_factor
+
+                second_pass_kwargs = {
+                    "prompt": prompt, "negative_prompt": negative_prompt, "height": second_pass_height, "width": second_pass_width,
+                    "num_frames": chunk.shape[2], "latents": chunk, "guidance_scale": float(guidance_scale),
+                    "output_type": "latent", "generator": torch.Generator(device=self.device).manual_seed(used_seed),
+                    **(self.config.get("second_pass", {}))
+                }
+                
+                refined_chunk = self.pipeline(**second_pass_kwargs).images
+                final_latents_list.append(refined_chunk.detach().clone())
+            
+            del upsampled_latents, chunk1, chunk2; torch.cuda.empty_cache()
+
+            final_latents = torch.cat(final_latents_list, dim=2)
+            log_tensor_info(final_latents, "Latentes Upscaled/Refinados Finais")
+
+            latents_cpu = final_latents.detach().to("cpu")
+            tensor_path = os.path.join(results_dir, f"latents_refined_{used_seed}.pt")
+            torch.save(latents_cpu, tensor_path)
+            
+            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, "refined_video", 24.0, temp_dir, results_dir, used_seed)
+            del pixel_tensor, final_latents
+
+        self._log_gpu_memory("Fim do Upscale/Denoise")
+        return video_path, tensor_path
+
+    def encode_mp4(self, latents_path: str, fps: int = 24):
+        print("\n--- INICIANDO ETAPA 3: DECODIFICAÇÃO FINAL ---")
+        self._log_gpu_memory("Início do Encode MP4")
+
+        latents = torch.load(latents_path)
+        seed = random.randint(0, 99999)
+        temp_dir = tempfile.mkdtemp(prefix="ltxv_enc_"); self._register_tmp_dir(temp_dir)
+        results_dir = "/app/output"; os.makedirs(results_dir, exist_ok=True)
+        
+        total_frames = latents.shape[2]
+        mid_point = total_frames // 2
+        chunk1_latents = latents[:, :, :mid_point, :, :]
+        chunk2_latents = latents[:, :, mid_point:, :, :]
+        
+        video_parts = []
+        with torch.autocast(device_type="cuda", dtype=self.runtime_autocast_dtype, enabled=self.device == 'cuda'):
+            for i, chunk in enumerate([chunk1_latents, chunk2_latents]):
+                if chunk.shape[2] == 0: continue
+                print(f"  - Decodificando chunk {i+1}/{2}")
+                pixel_chunk = vae_manager_singleton.decode(chunk.to(self.device), decode_timestep=float(self.config.get("decode_timestep", 0.05)))
+                
+                part_path = os.path.join(temp_dir, f"part_{i}.mp4")
+                video_encode_tool_singleton.save_video_from_tensor(pixel_chunk, part_path, fps=fps)
+                video_parts.append(part_path)
+                del pixel_chunk; torch.cuda.empty_cache()
+        
+        final_video_path = os.path.join(results_dir, f"final_concatenated_{seed}.mp4")
+        self._concat_mp4s_no_reencode(video_parts, final_video_path)
+        
+        print(f"Encode final concluído: {final_video_path}")
+        self._log_gpu_memory("Fim do Encode MP4")
+        return final_video_path
+
+# --- INSTANCIAÇÃO DO SERVIÇO ---
+print("Criando instância do VideoService. O carregamento do modelo começará agora...")
+video_generation_service = VideoService()
+print("Instância do VideoService pronta para uso.")