-
Notifications
You must be signed in to change notification settings - Fork 4
/
app.py
33 lines (27 loc) · 1.33 KB
/
app.py
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
from io import BytesIO
import base64
import torch
from diffusers import StableVideoDiffusionPipeline
from diffusers.utils import load_image, export_to_video
class InferlessPythonModel:
def initialize(self):
self.pipe = StableVideoDiffusionPipeline.from_pretrained("stabilityai/stable-video-diffusion-img2vid", torch_dtype=torch.float16, variant="fp16")
# self.pipe.enable_model_cpu_offload()
self.pipe.to("cuda")
self.pipe.unet = torch.compile(self.pipe.unet, mode="reduce-overhead", fullgraph=True)
self.pipe.vae = torch.compile(self.pipe.vae, mode="reduce-overhead", fullgraph=True)
def infer(self,inputs):
image_url = inputs['image_url']
image = load_image(image_url)
image = image.resize((1024, 576))
generator = torch.manual_seed(42)
frames = self.pipe(image, decode_chunk_size=8, generator=generator).frames[0]
export_to_video(frames, "generated.mp4", fps=7)
with open("generated.mp4", "rb") as video_file:
video_binary_data = video_file.read()
video_bytes_io = BytesIO(video_binary_data)
base64_encoded_data = base64.b64encode(video_bytes_io.read())
base64_string = base64_encoded_data.decode("utf-8")
return {"generated_video": base64_string}
def finalize(self):
pass