Update api/ltx_server_refactored.py

api/ltx_server_refactored.py

@@ -211,21 +211,11 @@ def create_ltx_video_pipeline(
         prompt_enhancer_llm_model = None
         prompt_enhancer_llm_tokenizer = None
 
-    # --- Ajuste global de precisão coerente ---
-    if precision in ["float8_e4m3fn", "bfloat16"]:
-        dtype_target = torch.bfloat16
-    elif precision == "mixed_precision":
-        dtype_target = torch.float16
-    else:
-        dtype_target = torch.float32
-    
-    for m in [vae, transformer, text_encoder]:
-        m.to(dtype_target)
-        
-    # garante coerência geral da pipeline
-    pipeline_dtype = dtype_target
+    vae = vae.to(torch.bfloat16)
+    if precision == "bfloat16" and transformer.dtype != torch.bfloat16:
+        transformer = transformer.to(torch.bfloat16)
+    text_encoder = text_encoder.to(torch.bfloat16)
 
-    
     # Use submodels for the pipeline
     submodel_dict = {
         "transformer": transformer,
@@ -364,48 +354,18 @@ class VideoService:
             "guidance_scale": float(guidance_scale),
             **(self.config.get("first_pass", {}))
         }
-        
+
         temp_dir = tempfile.mkdtemp(prefix="ltxv_low_")
         self._register_tmp_dir(temp_dir)
         
         try:
-            #with torch.autocast(device_type=self.device.split(':')[0], dtype=self.runtime_autocast_dtype, enabled=(self.device == 'cuda')):
-            latents_low = self.pipeline(**first_pass_kwargs).images
-
-            #with torch.autocast(device_type=self.device.split(':')[0], dtype=self.runtime_autocast_dtype, enabled=(self.device == 'cuda')):
-            #    latents = self.pipeline(**first_pass_kwargs).images
-            #    pixel_tensor = vae_manager_singleton.decode(latents.clone(), decode_timestep=float(self.config.get("decode_timestep", 0.05)))
-            #    video_path = self._save_video_from_tensor(pixel_tensor, "low_res_video", used_seed, temp_dir)
-            #    latents_path = self._save_latents_to_disk(latents, "latents_low_res", used_seed)
-
-            # ETAPA 2: Upscale dos latentes
-            #with torch.autocast(device_type=self.device.split(':')[0], dtype=self.runtime_autocast_dtype, enabled=(self.device == 'cuda')):
-            #upscaled_height, upscaled_width = downscaled_height * 2, downscaled_width * 2
-            #upscaled_latents = self._upsample_and_filter_latents(latents_low)
+            with torch.autocast(device_type=self.device.split(':')[0], dtype=self.runtime_autocast_dtype, enabled=(self.device == 'cuda')):
+                latents = self.pipeline(**first_pass_kwargs).images
+                pixel_tensor = vae_manager_singleton.decode(latents.clone(), decode_timestep=float(self.config.get("decode_timestep", 0.05)))
+                video_path = self._save_video_from_tensor(pixel_tensor, "low_res_video", used_seed, temp_dir)
+                latents_path = self._save_latents_to_disk(latents, "latents_low_res", used_seed)
                 
-            # ETAPA 3: Denoise final em alta resolução
-            second_pass_kwargs = {
-                "prompt": prompt,
-                "negative_prompt": negative_prompt,
-                "height": downscaled_height,
-                "width": downscaled_width,
-                "num_frames": max(3, (actual_num_frames//8)*8)+1,                "frame_rate": int(DEFAULT_FPS),
-                "generator": torch.Generator(device=self.device).manual_seed(used_seed),
-                "output_type": "latent",
-                "is_video": True, 
-                "latents": latents_low,
-                "vae_per_channel_normalize": True,
-                "conditioning_items": conditioning_items,
-                "guidance_scale": float(guidance_scale),
-                **(self.config.get("second_pass", {}))
-            }
-
-            #with torch.autocast(device_type=self.device.split(':')[0], dtype=self.runtime_autocast_dtype, enabled=(self.device == 'cuda')):
-            #latents = self.pipeline(**second_pass_kwargs).images
-            pixel_tensor = vae_manager_singleton.decode(latents_low.clone(), decode_timestep=float(self.config.get("decode_timestep", 0.05)))
-            video_path = self._save_video_from_tensor(pixel_tensor, "low_res_video", used_seed, temp_dir)
-            latents_path = self._save_latents_to_disk(latents, "latents_low_res", used_seed)
-            
+      
             return video_path, latents_path, used_seed
         
         except Exception as e: