| 44 | |
| 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, |
| 91 | motion_video: Path = Input( |
| 92 | description="Reference video file containing the motion to be mimicked" |
| 93 | ), |
| 94 | appearance_image: Path = Input( |
| 95 | description="Reference image file for the appearance of the generated video" |
| 96 | ), |
| 97 | resolution: int = Input( |
| 98 | description="Height of the output video in pixels. Width is automatically calculated.", |
| 99 | default=576, |
| 100 | ge=64, |
| 101 | le=1024, |
| 102 | ), |
| 103 | chunk_size: int = Input( |
nothing calls this directly
no outgoing calls
no test coverage detected