Load the model into memory to make running multiple predictions efficient
(self)
| 45 | |
| 46 | class Predictor(BasePredictor): |
| 47 | def setup(self): |
| 48 | """Load the model into memory to make running multiple predictions efficient""" |
| 49 | |
| 50 | if not os.path.exists(MODEL_CACHE): |
| 51 | os.makedirs(MODEL_CACHE) |
| 52 | model_files = [ |
| 53 | "DWPose.tar", |
| 54 | "MimicMotion.pth", |
| 55 | "MimicMotion_1-1.pth", |
| 56 | "SVD.tar", |
| 57 | ] |
| 58 | for model_file in model_files: |
| 59 | url = BASE_URL + model_file |
| 60 | filename = url.split("/")[-1] |
| 61 | dest_path = os.path.join(MODEL_CACHE, filename) |
| 62 | if not os.path.exists(dest_path.replace(".tar", "")): |
| 63 | download_weights(url, dest_path) |
| 64 | |
| 65 | self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| 66 | print(f"Using device: {self.device}") |
| 67 | |
| 68 | # Move imports here and make them global |
| 69 | # This ensures model files are downloaded before importing mimicmotion modules |
| 70 | global MimicMotionPipeline, create_pipeline, save_to_mp4, get_video_pose, get_image_pose |
| 71 | from mimicmotion.pipelines.pipeline_mimicmotion import MimicMotionPipeline |
| 72 | from mimicmotion.utils.loader import create_pipeline |
| 73 | from mimicmotion.utils.utils import save_to_mp4 |
| 74 | from mimicmotion.dwpose.preprocess import get_video_pose, get_image_pose |
| 75 | |
| 76 | # Load config with new checkpoint as default |
| 77 | self.config = OmegaConf.create( |
| 78 | { |
| 79 | "base_model_path": "models/SVD/stable-video-diffusion-img2vid-xt-1-1", |
| 80 | "ckpt_path": "models/MimicMotion_1-1.pth", |
| 81 | } |
| 82 | ) |
| 83 | |
| 84 | # Create the pipeline with the new checkpoint |
| 85 | self.pipeline = create_pipeline(self.config, self.device) |
| 86 | self.current_checkpoint = "v1-1" |
| 87 | self.current_dtype = torch.get_default_dtype() |
| 88 | |
| 89 | def predict( |
| 90 | self, |
nothing calls this directly
no test coverage detected