Update api/ltx_server_refactored_complete.py

api/ltx_server_refactored_complete.py

@@ -54,9 +54,16 @@ add_deps_to_path()
 
 # --- PROJECT IMPORTS ---
 try:
+    from ltx_video.pipelines.pipeline_ltx_video import LTXVideoPipeline, create_latent_upsampler # E outros...
+    from ltx_video.models.autoencoders.causal_video_autoencoder import CausalVideoAutoencoder
+    from ltx_video.models.transformers.transformer3d import Transformer3DModel
+    from ltx_video.models.transformers.symmetric_patchifier import SymmetricPatchifier
+    from ltx_video.schedulers.rf import RectifiedFlowScheduler
+    from transformers import T5EncoderModel, T5Tokenizer
+    from safetensors import safe_open
     from api.gpu_manager import gpu_manager
     from ltx_video.models.autoencoders.vae_encode import (normalize_latents, un_normalize_latents)
-    from ltx_video.pipelines.pipeline_ltx_video import (ConditioningItem, LTXMultiScalePipeline, adain_filter_latent, create_latent_upsampler, create_ltx_video_pipeline)
+    from ltx_video.pipelines.pipeline_ltx_video import (ConditioningItem, LTXMultiScalePipeline, adain_filter_latent, create_latent_upsampler)
     from ltx_video.utils.inference_utils import load_image_to_tensor_with_resize_and_crop
     from managers.vae_manager import vae_manager_singleton
     from tools.video_encode_tool import video_encode_tool_singleton
@@ -158,28 +165,82 @@ class VideoService:
         with open(config_path, "r") as file:
             return yaml.safe_load(file)
 
-    def _load_models(self) -> Tuple[LTXMultiScalePipeline, Optional[torch.nn.Module]]:
-        """Loads models from cache to CPU."""
+    def _load_models(self) -> Tuple[LTXVideoPipeline, Optional[torch.nn.Module]]:
+        """
+        Carrega todos os sub-modelos do pipeline na CPU.
+        Esta função substitui a necessidade de chamar a `create_ltx_video_pipeline` externa,
+        dando-nos controle total sobre o processo.
+        """
         t0 = time.perf_counter()
-        logging.info("Loading LTX models from cache to CPU...")
-        
-        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,
-        )
+        logging.info("Carregando sub-modelos do LTX para a CPU...")
+
+        ckpt_path = Path(self.config["checkpoint_path"])
+        if not ckpt_path.is_file():
+            raise FileNotFoundError(f"Arquivo de checkpoint principal não encontrado em: {ckpt_path}")
+
+        # 1. Carrega Metadados do Checkpoint
+        with safe_open(ckpt_path, framework="pt") as f:
+            metadata = f.metadata() or {}
+            config_str = metadata.get("config", "{}")
+            configs = json.loads(config_str)
+            allowed_inference_steps = configs.get("allowed_inference_steps")
+
+        # 2. Carrega os Componentes Individuais (todos na CPU)
+        #    O `.from_pretrained(ckpt_path)` é inteligente e carrega os pesos corretos do arquivo .safetensors.
+        logging.info("Carregando VAE...")
+        vae = CausalVideoAutoencoder.from_pretrained(ckpt_path).to("cpu")
+
+        logging.info("Carregando Transformer...")
+        transformer = Transformer3DModel.from_pretrained(ckpt_path).to("cpu")
+
+        logging.info("Carregando Scheduler...")
+        scheduler = RectifiedFlowScheduler.from_pretrained(ckpt_path)
+
+        logging.info("Carregando Text Encoder e Tokenizer...")
+        text_encoder_path = self.config["text_encoder_model_name_or_path"]
+        text_encoder = T5EncoderModel.from_pretrained(text_encoder_path, subfolder="text_encoder").to("cpu")
+        tokenizer = T5Tokenizer.from_pretrained(text_encoder_path, subfolder="tokenizer")
+
+        patchifier = SymmetricPatchifier(patch_size=1)
+
+        # 3. Define a precisão dos modelos (ainda na CPU, será aplicado na GPU depois)
+        precision = self.config.get("precision", "bfloat16")
+        if precision == "bfloat16":
+            vae.to(torch.bfloat16)
+            transformer.to(torch.bfloat16)
+            text_encoder.to(torch.bfloat16)
         
+        # 4. Monta o objeto do Pipeline com os componentes carregados
+        logging.info("Montando o objeto LTXVideoPipeline...")
+        submodel_dict = {
+            "transformer": transformer,
+            "patchifier": patchifier,
+            "text_encoder": text_encoder,
+            "tokenizer": tokenizer,
+            "scheduler": scheduler,
+            "vae": vae,
+            "allowed_inference_steps": allowed_inference_steps,
+            # Os prompt enhancers são opcionais e não são carregados por padrão para economizar memória
+            "prompt_enhancer_image_caption_model": None,
+            "prompt_enhancer_image_caption_processor": None,
+            "prompt_enhancer_llm_model": None,
+            "prompt_enhancer_llm_tokenizer": None,
+        }
+        pipeline = LTXVideoPipeline(**submodel_dict)
+
+        # 5. Carrega o Latent Upsampler (também na CPU)
         latent_upsampler = None
         if self.config.get("spatial_upscaler_model_path"):
+            logging.info("Carregando Latent Upsampler...")
             spatial_path = self.config["spatial_upscaler_model_path"]
             latent_upsampler = create_latent_upsampler(spatial_path, device="cpu")
+            if precision == "bfloat16":
+                latent_upsampler.to(torch.bfloat16)
 
-        logging.info(f"Models loaded on CPU in {time.perf_counter()-t0:.2f}s")
+        logging.info(f"Modelos LTX carregados na CPU em {time.perf_counter()-t0:.2f}s")
         return pipeline, latent_upsampler
-
+        
+    
     def move_to_device(self, main_device_str: str, vae_device_str: str):
         """Moves pipeline components to their target devices."""
         target_main_device = torch.device(main_device_str)