Update app.py

app.py

@@ -15,8 +15,6 @@ import cv2
 import shutil
 import glob
 from pathlib import Path
-from diffusers import AutoModel
-from diffusers.hooks import apply_group_offloading
 
 import warnings
 import logging
@@ -37,55 +35,34 @@ dtype = torch.bfloat16
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 
-# 1. Definir o repositório base
-base_model_repo = "Lightricks/LTX-Video"
+single_file_url = "https://huggingface.co/Lightricks/LTX-Video/resolve/main/ltxv-13b-0.9.8-distilled-fp8.safetensors"
 
-# 2. Carregar o Transformer separadamente para aplicar o casting FP8
-print("Carregando Transformer para otimização FP8...")
-transformer = AutoModel.from_pretrained(
-    base_model_repo,
-    subfolder="transformer",
-    torch_dtype=dtype
-)
-# Habilita a conversão dinâmica para FP8 (requer hardware compatível para funcionar)
-print("Habilitando layerwise casting para FP8...")
-transformer.enable_layerwise_casting(
-    storage_dtype=torch.float8_e4m3fn, compute_dtype=dtype
-)
-
-print("Desativando 'dynamic shifting' para compatibilidade com a pipeline.")
-transformer.config.use_dynamic_shifting = False
-
-# 3. Carregar a pipeline completa, injetando o Transformer já otimizado
-print(f"Carregando a arquitetura da pipeline de {base_model_repo}...")
-pipeline = LTXConditionPipeline.from_pretrained(
-    base_model_repo,
-    transformer=transformer, # Injeta o transformer otimizado
-    torch_dtype=dtype,
+pipeline = LTXConditionPipeline.from_single_file(
+    single_file_url,
+    offload_state_dict=False,
+    dtype=torch.bfloat16, # Use o dtype apropriado. Para FP8, pode ser necessário torch.float8_e4m3fn.
     cache_dir=os.getenv("HF_HOME_CACHE"),
     token=os.getenv("HF_TOKEN"),
 )
 
-# 4. Carregar o upsampler (seu repositório é separado e está correto)
-print("Carregando upsampler...")
+# Carregamento das pipelines
+#pipeline = LTXConditionPipeline.from_pretrained(
+#    "Lightricks/LTX-Video-0.9.8-13B-distilled",
+#    offload_state_dict=False,
+#    torch_dtype=torch.bfloat16,
+#    cache_dir=os.getenv("HF_HOME_CACHE"),
+#    token=os.getenv("HF_TOKEN"),
+#)
+
 pipe_upsample = LTXLatentUpsamplePipeline.from_pretrained(
     "Lightricks/ltxv-spatial-upscaler-0.9.7",
     cache_dir=os.getenv("HF_HOME_CACHE"),
-    vae=pipeline.vae,
-    torch_dtype=dtype
+    vae=pipeline.vae, torch_dtype=dtype
 )
 
-
-# 5. Aplicar o descarregamento de grupos para economizar VRAM
-print("Aplicando otimizações de group-offloading para economizar VRAM...")
-onload_device = torch.device("cuda")
-offload_device = torch.device("cpu")
-# O Transformer já tem um método integrado
-pipeline.transformer.enable_group_offload(onload_device=onload_device, offload_device=offload_device, offload_type="leaf_level", use_stream=True)
-# Para os outros componentes, usamos a função auxiliar
-apply_group_offloading(pipeline.text_encoder, onload_device=onload_device, offload_type="block_level", num_blocks_per_group=2)
-apply_group_offloading(pipeline.vae, onload_device=onload_device, offload_type="leaf_level")
-
+pipeline.to(device)
+pipe_upsample.to(device)
+pipeline.vae.enable_tiling()
 
 
 current_dir = Path(__file__).parent