fix: clean up

This commit is contained in:
mbukeRepo 2024-07-07 15:46:28 +02:00
parent 0b303fc39c
commit eec7aa3337
2 changed files with 12 additions and 6 deletions

View File

@ -29,4 +29,5 @@ build:
- "matplotlib==3.9.0" - "matplotlib==3.9.0"
- "imageio-ffmpeg==0.5.1" - "imageio-ffmpeg==0.5.1"
- "tyro==0.8.5" - "tyro==0.8.5"
- "gradio==3.48.0"
predict: "predict.py:Predictor" predict: "predict.py:Predictor"

View File

@ -1,9 +1,10 @@
from cog import BasePredictor, Input, Path from cog import BasePredictor, Input, Path, File
from src.config.argument_config import ArgumentConfig from src.config.argument_config import ArgumentConfig
from src.config.inference_config import InferenceConfig from src.config.inference_config import InferenceConfig
from src.config.crop_config import CropConfig from src.config.crop_config import CropConfig
from src.live_portrait_pipeline import LivePortraitPipeline from src.live_portrait_pipeline import LivePortraitPipeline
import requests
class Predictor(BasePredictor): class Predictor(BasePredictor):
@ -16,14 +17,18 @@ class Predictor(BasePredictor):
def predict( def predict(
self, self,
image: Path = Input(description="Portrait image"), image: Path = Input(description="Portrait image")
driving_info: Path = Input(
description="driving video or template (.pkl format)"
),
) -> Path: ) -> Path:
"""Run a single prediction on the model""" """Run a single prediction on the model"""
video_path, _ = self.live_portrait_pipeline.execute( video_path, _ = self.live_portrait_pipeline.execute(
ArgumentConfig(source_image=image, driving_info=driving_info, output_dir="/tmp/") ArgumentConfig(
source_image=image,
driving_info="assets/examples/driving/d0.mp4",
output_dir="/tmp/",
flag_pasteback=False,
flag_do_crop=False,
)
) )
return Path(video_path) return Path(video_path)