Spaces:
Running
on
Zero
Running
on
Zero
File size: 2,028 Bytes
84d7244 d67e138 84d7244 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 |
from models.SpaTrackV2.models.predictor import Predictor
from models.SpaTrackV2.models.vggt4track.models.vggt_moe import VGGT4Track
import logging
from .config import (
TTM_COG_AVAILABLE, TTM_WAN_AVAILABLE,
TTM_COG_MODEL_ID, TTM_WAN_MODEL_ID, TTM_DTYPE
)
logger = logging.getLogger(__name__)
vggt4track_model = None
tracker_model = None
ttm_cog_pipeline = None
ttm_wan_pipeline = None
def init_spatial_models():
global vggt4track_model, tracker_model
print("🚀 Initializing models...")
vggt4track_model = VGGT4Track.from_pretrained(
"Yuxihenry/SpatialTrackerV2_Front")
vggt4track_model.eval()
vggt4track_model = vggt4track_model.to("cuda")
tracker_model = Predictor.from_pretrained(
"Yuxihenry/SpatialTrackerV2-Offline")
tracker_model.eval()
print("✅ Spatial Models loaded successfully!")
def get_ttm_cog_pipeline():
global ttm_cog_pipeline
if ttm_cog_pipeline is None and TTM_COG_AVAILABLE:
from diffusers import CogVideoXImageToVideoPipeline
logger.info("Loading TTM CogVideoX pipeline...")
ttm_cog_pipeline = CogVideoXImageToVideoPipeline.from_pretrained(
TTM_COG_MODEL_ID,
torch_dtype=TTM_DTYPE,
low_cpu_mem_usage=True,
)
ttm_cog_pipeline.vae.enable_tiling()
ttm_cog_pipeline.vae.enable_slicing()
logger.info("TTM CogVideoX pipeline loaded successfully!")
return ttm_cog_pipeline
def get_ttm_wan_pipeline():
global ttm_wan_pipeline
if ttm_wan_pipeline is None and TTM_WAN_AVAILABLE:
from diffusers import WanImageToVideoPipeline
logger.info("Loading TTM Wan 2.2 pipeline...")
ttm_wan_pipeline = WanImageToVideoPipeline.from_pretrained(
TTM_WAN_MODEL_ID,
torch_dtype=TTM_DTYPE,
)
ttm_wan_pipeline.vae.enable_tiling()
ttm_wan_pipeline.vae.enable_slicing()
logger.info("TTM Wan 2.2 pipeline loaded successfully!")
return ttm_wan_pipeline
init_spatial_models()
|