Spaces:
Configuration error
Configuration error
| import os | |
| from typing import List | |
| from cog import BasePredictor, Input, Path | |
| import subprocess | |
| import shutil | |
| MODEL_CACHE = "model-cache" | |
| class Predictor(BasePredictor): | |
| def setup(self): | |
| pass | |
| def predict( | |
| self, | |
| prompt: str = Input( | |
| description="Input prompt", default="An astronaut riding a horse" | |
| ), | |
| negative_prompt: str = Input( | |
| description="Negative prompt", default=None | |
| ), | |
| init_video: Path = Input( | |
| description="URL of the initial video (optional)", default=None | |
| ), | |
| init_weight: float = Input( | |
| description="Strength of init_video", default=0.5 | |
| ), | |
| num_frames: int = Input( | |
| description="Number of frames for the output video", default=24 | |
| ), | |
| num_inference_steps: int = Input( | |
| description="Number of denoising steps", ge=1, le=500, default=50 | |
| ), | |
| width: int = Input( | |
| description="Width of the output video", ge=256, default=576 | |
| ), | |
| height: int = Input( | |
| description="Height of the output video", ge=256, default=320 | |
| ), | |
| guidance_scale: float = Input( | |
| description="Guidance scale", ge=1.0, le=100.0, default=7.5 | |
| ), | |
| fps: int = Input(description="fps for the output video", default=8), | |
| model: str = Input( | |
| description="Model to use", default="xl", choices=["xl", "576w", "potat1", "animov-512x"] | |
| ), | |
| batch_size: int = Input(description="Batch size", default=1, ge=1), | |
| remove_watermark: bool = Input( | |
| description="Remove watermark", default=False | |
| ), | |
| seed: int = Input( | |
| description="Random seed. Leave blank to randomize the seed", default=None | |
| ), | |
| ) -> List[Path]: | |
| if seed is None: | |
| seed = int.from_bytes(os.urandom(2), "big") | |
| print(f"Using seed: {seed}") | |
| shutil.rmtree("output", ignore_errors=True) | |
| os.makedirs("output", exist_ok=True) | |
| args = { | |
| "prompt": prompt, | |
| "negative_prompt": negative_prompt, | |
| "batch_size": batch_size, | |
| "num_frames": num_frames, | |
| "num_steps": num_inference_steps, | |
| "seed": seed, | |
| "guidance-scale": guidance_scale, | |
| "width": width, | |
| "height": height, | |
| "fps": fps, | |
| "device": "cuda", | |
| "output_dir": "output", | |
| "remove-watermark": remove_watermark, | |
| } | |
| args['model'] = MODEL_CACHE + "/" + model | |
| if init_video is not None: | |
| # for some reason I need to copy the file to make it work | |
| if os.path.exists("input.mp4"): | |
| os.unlink("input.mp4") | |
| shutil.copy(init_video, "input.mp4") | |
| args["init-video"] = "input.mp4" | |
| args["init-weight"] = init_weight | |
| print("init video", os.stat("input.mp4").st_size) | |
| cmd = ["python", "inference.py"] | |
| for k, v in args.items(): | |
| if not v is None: | |
| cmd.append(f"--{k}") | |
| cmd.append(str(v)) | |
| subprocess.check_call(cmd) | |
| # outputs = inference.run(**args) | |
| outputs = [] | |
| for f in os.listdir("output"): | |
| if f.endswith(".mp4"): | |
| outputs.append(Path(os.path.join("output", f))) | |
| return outputs | |