File size: 1,764 Bytes
ade3957 |
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 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 |
import torch
from diffusers import DiffusionPipeline, StableDiffusionImageVariationPipeline
from PIL import Image
import numpy as np
import cv2
class BootyShakerPipeline:
def __init__(self):
self.txt2video_pipe = DiffusionPipeline.from_pretrained("ChromiumPlutoniumAI/BootyShakerAI")
self.img2video_pipe = StableDiffusionImageVariationPipeline.from_pretrained("ChromiumPlutoniumAI/BootyShakerAI")
self.device = "cuda" if torch.cuda.is_available() else "cpu"
def generate_from_text(self, prompt, num_frames=16, fps=8):
video = self.txt2video_pipe(
prompt,
num_inference_steps=50,
num_frames=num_frames
).frames
return self.frames_to_video(video, fps)
def generate_from_image(self, image, num_frames=16, fps=8):
if isinstance(image, str):
image = Image.open(image)
video = self.img2video_pipe(
image,
num_inference_steps=50,
num_frames=num_frames
).frames
return self.frames_to_video(video, fps)
def apply_alterations(self, video, style="bounce"):
styles = {
"bounce": self.bounce_effect,
"wave": self.wave_effect,
"shake": self.shake_effect
}
return styles[style](video)
@staticmethod
def frames_to_video(frames, fps):
output_file = "output.mp4"
writer = cv2.VideoWriter(
output_file,
cv2.VideoWriter_fourcc(*"mp4v"),
fps,
(frames[0].shape[1], frames[0].shape[0])
)
for frame in frames:
writer.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
writer.release()
return output_file
|