Update api/ltx_server_refactored.py

api/ltx_server_refactored.py

@@ -322,7 +322,12 @@ class VideoService:
             
         return conditioning_items
 
-    def generate_low_resolution(self, prompt: str, negative_prompt: str, height: int, width: int, duration_secs: float, guidance_scale: float, seed: Optional[int] = None, conditioning_items: Optional[List[ConditioningItem]] = None) -> Tuple[str, str, int]:
+    def generate_low_resolution(
+        self, prompt: str, negative_prompt: str, 
+        height: int, width: int, duration_secs: float, 
+        guidance_scale: float, seed: Optional[int] = None, 
+        conditioning_items: Optional[List[ConditioningItem]] = None
+    ) -> Tuple[str, str, int]:
         """
         Gera um vídeo de baixa resolução e retorna os caminhos para o vídeo e os latentes.
         """
@@ -334,11 +339,17 @@ class VideoService:
         downscaled_height, downscaled_width = self._calculate_downscaled_dims(height, width)
 
         first_pass_kwargs = {
-            "prompt": prompt, "negative_prompt": negative_prompt, "height": downscaled_height,
-            "width": downscaled_width, "num_frames": actual_num_frames, "frame_rate": int(DEFAULT_FPS),
+            "prompt": prompt,
+            "negative_prompt": negative_prompt,
+            "height": downscaled_height,
+            "width": downscaled_width,
+            "num_frames": num_frames,
+            "frame_rate": int(DEFAULT_FPS),
             "generator": torch.Generator(device=self.device).manual_seed(used_seed),
-            "output_type": "latent", "conditioning_items": conditioning_items,
-            "guidance_scale": float(guidance_scale), **(self.config.get("first_pass", {}))
+            "output_type": "latent",
+            "conditioning_items": conditioning_items,
+            "guidance_scale": float(guidance_scale),
+            **(self.config.get("first_pass", {}))
         }
 
         temp_dir = tempfile.mkdtemp(prefix="ltxv_low_")
@@ -355,10 +366,11 @@ class VideoService:
                     latents_path=latents_path,
                     prompt=prompt,
                     negative_prompt=negative_prompt,
-                    height=downscaled_height, width=downscaled_width, num_frames=actual_num_frames, 
+                    height=downscaled_height, width=downscaled_width, 
+                    num_frames=actual_num_frames, 
                     guidance_scale=guidance_scale,
                     seed=used_seed,
-                    conditioning_items=conditioning_items
+                    conditioning_items=conditioning_items,
                 )
                 
                 print(f"[SUCCESS] PASSO 2 concluído. Vídeo final em: {final_video_path}")
@@ -372,7 +384,12 @@ class VideoService:
         finally:
             self._finalize()
 
-    def generate_upscale_denoise(self, latents_path: str, prompt: str, negative_prompt: str, height: int, width: int, num_frames: float, guidance_scale: float, seed: Optional[int] = None, conditioning_items: Optional[List[ConditioningItem]] = None) -> Tuple[str, str]:
+    def generate_upscale_denoise(
+        self, latents_path: str, prompt: str, 
+        negative_prompt: str, height: int, width: int, 
+        num_frames: float, guidance_scale: float, seed: Optional[int] = None, 
+        conditioning_items: Optional[List[ConditioningItem]] = None
+    ) -> Tuple[str, str]:
         """
         Aplica upscale, AdaIN e Denoise em latentes de baixa resolução usando um processo de chunking.
         """
@@ -386,8 +403,7 @@ class VideoService:
             latents_low = torch.load(latents_path).to(self.device)
             with torch.autocast(device_type=self.device.split(':')[0], dtype=self.runtime_autocast_dtype, enabled=(self.device == 'cuda')):
                 upsampled_latents = latents_low #self._upsample_and_filter_latents(latents_low)
-                del latents_low; torch.cuda.empty_cache()
-
+                
                 #chunks = self._split_latents_with_overlap(upsampled_latents)
                 #refined_chunks = []
 
@@ -395,20 +411,28 @@ class VideoService:
                     #if chunk.shape[2] <= 1: continue  # Pula chunks inválidos
 
                 chunk = upsampled_latents
-                if True:
-                    second_pass_height = chunk.shape[3] * self.pipeline.vae_scale_factor
-                    second_pass_width = chunk.shape[4] * self.pipeline.vae_scale_factor
+                
+                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),
-                        conditioning_items: conditioning_items
-                        **(self.config.get("second_pass", {}))
-                    }
-                    refined_chunk = self.pipeline(**second_pass_kwargs).images
-                    #refined_chunks.append(refined_chunk)
+                second_pass_kwargs = {
+                    "prompt": prompt,
+                    "negative_prompt": negative_prompt,
+                    "height": second_pass_height,
+                    "width": second_pass_width,
+                    "frame_rate": int(DEFAULT_FPS),
+                    "num_frames": num_frames,
+                    "latents": chunk,  # O tensor completo é passado aqui
+                    "guidance_scale": float(guidance_scale),
+                    "output_type": "latent",
+                    "generator": torch.Generator(device=self.device).manual_seed(used_seed),
+                    "conditioning_items": conditioning_items,
+                    **(self.config.get("second_pass", {}))
+                }
+                refined_chunk = self.pipeline(**second_pass_kwargs).images
+                #refined_chunks.append(refined_chunk)
+            
+            del latents_low; torch.cuda.empty_cache()
             
             final_latents = refined_chunk #self._merge_chunks_with_overlap(refined_chunks)
             #if LTXV_DEBUG: