videopython 0.1.3__tar.gz → 0.1.41__tar.gz

{videopython-0.1.3 → videopython-0.1.41}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: videopython
-Version: 0.1.3
+Version: 0.1.41
 Summary: Minimal video generation and processing library.
 Author-email: Bartosz Wójtowicz <bartoszwojtowicz@outlook.com>, Bartosz Rudnikowicz <bartoszrudnikowicz840@gmail.com>, Piotr Pukisz <piotr.pukisz@gmail.com>
 License:                                  Apache License
@@ -199,7 +199,7 @@ License:                                  Apache License
 Project-URL: Homepage, https://github.com/bartwojtowicz/videopython/
 Project-URL: Bug Reports, https://github.com/bartwojtowicz/videopython/issues
 Project-URL: Source, https://github.com/bartwojtowicz/videopython/
-Keywords: videopython,video,movie,opencv,generation,editing
+Keywords: python,videopython,video,movie,opencv,generation,editing
 Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.10
@@ -211,13 +211,22 @@ License-File: LICENSE
 Requires-Dist: click>=8.1.7
 Requires-Dist: numpy>=1.25.2
 Requires-Dist: opencv-python>=4.9.0.80
-Requires-Dist: pytest>=7.4.0
-Requires-Dist: transformers>=4.38.1
-Requires-Dist: diffusers>=0.26.3
-Requires-Dist: torch>=2.1.0
-Requires-Dist: stability-sdk>=0.8.5
-Requires-Dist: openai==1.3.5
+Requires-Dist: pillow>=10.3.0
 Requires-Dist: pydub>=0.25.1
+Requires-Dist: tqdm>=4.66.3
+Provides-Extra: dev
+Requires-Dist: black==24.3.0; extra == "dev"
+Requires-Dist: isort==5.12.0; extra == "dev"
+Requires-Dist: mypy==1.8.0; extra == "dev"
+Requires-Dist: pytest==7.4.0; extra == "dev"
+Requires-Dist: types-Pillow==10.2.0.20240213; extra == "dev"
+Requires-Dist: types-tqdm==4.66.0.20240106; extra == "dev"
+Requires-Dist: pydub-stubs==0.25.1.1; extra == "dev"
+Provides-Extra: generation
+Requires-Dist: accelerate>=0.29.2; extra == "generation"
+Requires-Dist: diffusers>=0.26.3; extra == "generation"
+Requires-Dist: torch>=2.1.0; extra == "generation"
+Requires-Dist: transformers>=4.38.1; extra == "generation"
 
 # About
 
@@ -235,41 +244,40 @@ sudo apt-get install ffmpeg
 
 ### Install with pip
 ```bash
-pip install videopython
+pip install videopython[generation]
 ```
+> You can install without `[generation]` dependencies for basic video handling and processing. 
+> The funcionalities found in `videopython.generation` won't work.
 
 ## Basic Usage
+> Using Nvidia A40 or better is recommended for the `videopython.generation` module.
 
 ```python
-from videopython.base.video import Video
-from videopython.base.transitions import FadeTransition
+# Generate image and animate it
+from videopython.generation import ImageToVideo
+from videopython.generation import TextToImage
 
-# Load video
-video = Video.from_path("tests/test_data/fast_benchmark.mp4")
-print(video.metadata)
-print(video.frames.shape) # Video is based on numpy representation of frames
+image = TextToImage().generate_image(prompt="Golden Retriever playing in the park")
+video = ImageToVideo().generate_video(image=image, fps=24)
 
-# Generate videos
-video1 = Video.from_prompt("Dogs playing in the snow.")
-video2 = Video.from_prompt("Dogs going back home.")
+# Video generation directly from prompt
+from videopython.generation import TextToVideo
+video_gen = TextToVideo()
+video = video_gen.generate_video("Dogs playing in the snow")
+for _ in range(10):
+    video += video_gen.generate_video("Dogs playing in the snow")
 
-# Add videos
-combined_video = video1 + video2
-print(combined_video.metadata)
+# Cut the first 2 seconds
+from videopython.base.transforms import CutSeconds
+transformed_video = CutSeconds(start_second=0, end_second=2).apply(video.copy())
 
-# Apply fade transition between videos
-fade = FadeTransition(0.5) # 0.5s effect time
-faded_video = fade.apply(videos=(video1, video2))
-print(faded_video.metadata)
+# Upsample to 30 FPS
+from videopython.base.transforms import ResampleFPS
+transformed_video = ResampleFPS(new_fps=30).apply(transformed_video)
 
-# Add audio from file
-faded_video.add_audio_from_file("tests/test_data/test_audio.mp3")
+# Resize to 1000x1000
+from videopython.base.transforms import Resize
+transformed_video = Resize(width=1000, height=1000).apply(transformed_video)
 
-# Save to a file
-faded_video.save("my_video.mp4")
-```
-
-### Running Unit Tests
-```bash
-PYTHONPATH=./src/ pytest
+filepath = transformed_video.save()
 ```

videopython-0.1.41/README.md

@@ -0,0 +1,53 @@
+# About
+
+Minimal video generation and processing library.
+
+## Setup 
+
+### Install ffmpeg
+```bash
+# Install with brew for MacOS:
+brew install ffmpeg
+# Install with apt-get for Ubuntu:
+sudo apt-get install ffmpeg
+```
+
+### Install with pip
+```bash
+pip install videopython[generation]
+```
+> You can install without `[generation]` dependencies for basic video handling and processing. 
+> The funcionalities found in `videopython.generation` won't work.
+
+## Basic Usage
+> Using Nvidia A40 or better is recommended for the `videopython.generation` module.
+
+```python
+# Generate image and animate it
+from videopython.generation import ImageToVideo
+from videopython.generation import TextToImage
+
+image = TextToImage().generate_image(prompt="Golden Retriever playing in the park")
+video = ImageToVideo().generate_video(image=image, fps=24)
+
+# Video generation directly from prompt
+from videopython.generation import TextToVideo
+video_gen = TextToVideo()
+video = video_gen.generate_video("Dogs playing in the snow")
+for _ in range(10):
+    video += video_gen.generate_video("Dogs playing in the snow")
+
+# Cut the first 2 seconds
+from videopython.base.transforms import CutSeconds
+transformed_video = CutSeconds(start_second=0, end_second=2).apply(video.copy())
+
+# Upsample to 30 FPS
+from videopython.base.transforms import ResampleFPS
+transformed_video = ResampleFPS(new_fps=30).apply(transformed_video)
+
+# Resize to 1000x1000
+from videopython.base.transforms import Resize
+transformed_video = Resize(width=1000, height=1000).apply(transformed_video)
+
+filepath = transformed_video.save()
+```

{videopython-0.1.3 → videopython-0.1.41}/pyproject.toml

@@ -1,5 +1,5 @@
 [build-system]
-requires = ["setuptools>=61.0"]
+requires = ["setuptools>=66.1"]
 build-backend = "setuptools.build_meta"
 
 [tool.setuptools.packages.find]
@@ -11,12 +11,13 @@ include = ["videopython.*"]
 
 [project]
 name = "videopython"
-version = "0.1.3"
+version = "0.1.41"
 description = "Minimal video generation and processing library."
 readme = "README.md"
 requires-python = ">=3.10"
 license = {file = "LICENSE"}
-keywords = ["videopython", "video", "movie", "opencv", "generation", "editing"]
+keywords = ["python", "videopython", "video", "movie", "opencv", "generation", "editing"]
+dynamic = ["dependencies", "optional-dependencies"]
 
 authors = [
   {name = "Bartosz Wójtowicz", email = "bartoszwojtowicz@outlook.com" },
@@ -32,20 +33,11 @@ classifiers = [
   "Operating System :: OS Independent",
 ]
 
-dependencies = [
-    "click>=8.1.7",
-    "numpy>=1.25.2",
-    "opencv-python>=4.9.0.80",
-    "pytest>=7.4.0",
-    "transformers>=4.38.1",
-    "diffusers>=0.26.3",
-    "torch>=2.1.0",
-    "stability-sdk>=0.8.5",
-    "openai==1.3.5",
-    "pydub>=0.25.1"
-]
+[tool.setuptools.dynamic]
+dependencies = {file = ["requirements.txt"]}
+optional-dependencies = { dev = {file = ["requirements-dev.txt"]}, generation = {file = ["requirements-generation.txt"]} }
 
-[project.urls]  # Optional
+[project.urls]
 "Homepage" = "https://github.com/bartwojtowicz/videopython/"
 "Bug Reports" = "https://github.com/bartwojtowicz/videopython/issues"
 "Source" = "https://github.com/bartwojtowicz/videopython/"

videopython-0.1.41/requirements-dev.txt

@@ -0,0 +1,7 @@
+black==24.3.0
+isort==5.12.0
+mypy==1.8.0
+pytest==7.4.0
+types-Pillow==10.2.0.20240213
+types-tqdm==4.66.0.20240106
+pydub-stubs==0.25.1.1

videopython-0.1.41/requirements-generation.txt

@@ -0,0 +1,4 @@
+accelerate>=0.29.2
+diffusers>=0.26.3
+torch>=2.1.0
+transformers>=4.38.1

videopython-0.1.41/requirements.txt

@@ -0,0 +1,6 @@
+click>=8.1.7
+numpy>=1.25.2
+opencv-python>=4.9.0.80
+pillow>=10.3.0
+pydub>=0.25.1
+tqdm>=4.66.3

videopython-0.1.41/src/videopython/base/effects.py

@@ -0,0 +1,183 @@
+from abc import ABC, abstractmethod
+from typing import Literal, final
+
+import cv2
+import numpy as np
+from PIL import Image
+from tqdm import tqdm
+
+from videopython.base.video import Video
+
+
+class Effect(ABC):
+    """Abstract class for effect on frames of video.
+
+    The effect must not change the number of frames and the shape of the frames.
+    """
+
+    @final
+    def apply(self, video: Video, start: float | None = None, stop: float | None = None) -> Video:
+        original_shape = video.video_shape
+        start = start if start is not None else 0
+        stop = stop if stop is not None else video.total_seconds
+        # Check for start and stop correctness
+        if not 0 <= start <= video.total_seconds:
+            raise ValueError(f"Video is only {video.total_seconds} long, but passed start: {start}!")
+        elif not start <= stop <= video.total_seconds:
+            raise ValueError(f"Video is only {video.total_seconds} long, but passed stop: {stop}!")
+        # Apply effect on video slice
+        effect_start_frame = round(start * video.fps)
+        effect_end_frame = round(stop * video.fps)
+        video_with_effect = self._apply(video[effect_start_frame:effect_end_frame])
+        old_audio = video.audio
+        video = Video.from_frames(
+            np.r_[
+                "0,2",
+                video.frames[:effect_start_frame],
+                video_with_effect.frames,
+                video.frames[effect_end_frame:],
+            ],
+            fps=video.fps,
+        )
+        video.audio = old_audio
+        # Check if dimensions didn't change
+        if not video.video_shape == original_shape:
+            raise RuntimeError("The effect must not change the number of frames and the shape of the frames!")
+
+        return video
+
+    @abstractmethod
+    def _apply(self, video: Video) -> Video:
+        pass
+
+
+class FullImageOverlay(Effect):
+    def __init__(self, overlay_image: np.ndarray, alpha: float | None = None, fade_time: float = 0.0):
+        if alpha is not None and not 0 <= alpha <= 1:
+            raise ValueError("Alpha must be in range [0, 1]!")
+        elif not (overlay_image.ndim == 3 and overlay_image.shape[-1] in [3, 4]):
+            raise ValueError("Only RGB and RGBA images are supported as an overlay!")
+        elif alpha is None:
+            alpha = 1.0
+
+        if overlay_image.shape[-1] == 3:
+            overlay_image = np.dstack([overlay_image, np.full(overlay_image.shape[:2], 255, dtype=np.uint8)])
+
+        self.alpha = alpha
+        self.overlay = overlay_image.astype(np.uint8)
+        self.fade_time = fade_time
+
+    def _overlay(self, img: np.ndarray, alpha: float = 1.0) -> np.ndarray:
+        img_pil = Image.fromarray(img)
+        overlay = self.overlay.copy()
+        overlay[:, :, 3] = overlay[:, :, 3] * (self.alpha * alpha)
+        overlay_pil = Image.fromarray(overlay)
+        img_pil.paste(overlay_pil, (0, 0), overlay_pil)
+        return np.array(img_pil)
+
+    def _apply(self, video: Video) -> Video:
+        if not video.frame_shape == self.overlay[:, :, :3].shape:
+            raise ValueError(
+                f"Mismatch of overlay shape `{self.overlay.shape}` with video shape: `{video.frame_shape}`!"
+            )
+        elif not (0 <= 2 * self.fade_time <= video.total_seconds):
+            raise ValueError(f"Video is only {video.total_seconds}s long, but fade time is {self.fade_time}s!")
+
+        print("Overlaying video...")
+        if self.fade_time == 0:
+            video.frames = np.array([self._overlay(frame) for frame in tqdm(video.frames)], dtype=np.uint8)
+        else:
+            num_video_frames = len(video.frames)
+            num_fade_frames = round(self.fade_time * video.fps)
+            new_frames = []
+            for i, frame in enumerate(tqdm(video.frames)):
+                frames_dist_from_end = min(i, num_video_frames - i)
+                if frames_dist_from_end >= num_fade_frames:
+                    fade_alpha = 1.0
+                else:
+                    fade_alpha = frames_dist_from_end / num_fade_frames
+                new_frames.append(self._overlay(frame, fade_alpha))
+            video.frames = np.array(new_frames, dtype=np.uint8)
+        return video
+
+
+class Blur(Effect):
+    def __init__(
+        self,
+        mode: Literal["constant", "ascending", "descending"],
+        iterations: int,
+        kernel_size: tuple[int, int] = (5, 5),
+    ):
+        if iterations < 1:
+            raise ValueError("Iterations must be at least 1!")
+        self.mode = mode
+        self.iterations = iterations
+        self.kernel_size = kernel_size
+
+    def _apply(self, video: Video) -> Video:
+        n_frames = len(video.frames)
+        new_frames = []
+        if self.mode == "constant":
+            for frame in video.frames:
+                blurred_frame = frame
+                for _ in range(self.iterations):
+                    blurred_frame = cv2.GaussianBlur(blurred_frame, self.kernel_size, 0)
+                new_frames.append(blurred_frame)
+        elif self.mode == "ascending":
+            for i, frame in tqdm(enumerate(video.frames)):
+                frame_iterations = max(1, round((i / n_frames) * self.iterations))
+                blurred_frame = frame
+                for _ in range(frame_iterations):
+                    blurred_frame = cv2.GaussianBlur(blurred_frame, self.kernel_size, 0)
+                new_frames.append(blurred_frame)
+        elif self.mode == "descending":
+            for i, frame in tqdm(enumerate(video.frames)):
+                frame_iterations = max(round(((n_frames - i) / n_frames) * self.iterations), 1)
+                blurred_frame = frame
+                for _ in range(frame_iterations):
+                    blurred_frame = cv2.GaussianBlur(blurred_frame, self.kernel_size, 0)
+                new_frames.append(blurred_frame)
+        else:
+            raise ValueError(f"Unknown mode: `{self.mode}`.")
+        video.frames = np.asarray(new_frames)
+        return video
+
+
+class Zoom(Effect):
+    def __init__(self, zoom_factor: float, mode: Literal["in", "out"]):
+        if zoom_factor <= 1:
+            raise ValueError("Zoom factor must be greater than 1!")
+        self.zoom_factor = zoom_factor
+        self.mode = mode
+
+    def _apply(self, video: Video) -> Video:
+        n_frames = len(video.frames)
+        new_frames = []
+
+        width = video.metadata.width
+        height = video.metadata.height
+        crop_sizes_w, crop_sizes_h = np.linspace(width // self.zoom_factor, width, n_frames), np.linspace(
+            height // self.zoom_factor, height, n_frames
+        )
+
+        if self.mode == "in":
+            for frame, w, h in tqdm(zip(video.frames, reversed(crop_sizes_w), reversed(crop_sizes_h))):
+
+                x = width / 2 - w / 2
+                y = height / 2 - h / 2
+
+                cropped_frame = frame[round(y) : round(y + h), round(x) : round(x + w)]
+                zoomed_frame = cv2.resize(cropped_frame, (width, height))
+                new_frames.append(zoomed_frame)
+        elif self.mode == "out":
+            for frame, w, h in tqdm(zip(video.frames, crop_sizes_w, crop_sizes_h)):
+                x = width / 2 - w / 2
+                y = height / 2 - h / 2
+
+                cropped_frame = frame[round(y) : round(y + h), round(x) : round(x + w)]
+                zoomed_frame = cv2.resize(cropped_frame, (width, height))
+                new_frames.append(zoomed_frame)
+        else:
+            raise ValueError(f"Unknown mode: `{self.mode}`.")
+        video.frames = np.asarray(new_frames)
+        return video

videopython-0.1.41/src/videopython/base/exceptions.py

@@ -0,0 +1,2 @@
+class OutOfBoundsError(Exception):
+    pass

{videopython-0.1.3 → videopython-0.1.41}/src/videopython/base/transforms.py

@@ -1,5 +1,7 @@
 from abc import ABC, abstractmethod
+from enum import Enum
 from multiprocessing import Pool
+from typing import Literal
 
 import cv2
 import numpy as np
@@ -71,9 +73,11 @@ class CutSeconds(Transformation):
 
 
 class Resize(Transformation):
-    def __init__(self, new_width: int, new_height: int):
-        self.new_width = new_width
-        self.new_height = new_height
+    def __init__(self, width: int | None = None, height: int | None = None):
+        self.width = width
+        self.height = height
+        if width is None and height is None:
+            raise ValueError("You must provide either `width` or `height`!")
 
     def _resize_frame(self, frame: np.ndarray, new_width: int, new_height: int) -> np.ndarray:
         return cv2.resize(
@@ -83,10 +87,25 @@ class Resize(Transformation):
         )
 
     def apply(self, video: Video) -> Video:
+        if self.width and self.height:
+            new_height = self.height
+            new_width = self.width
+        elif self.height is None and self.width:
+            video_height = video.video_shape[1]
+            video_width = video.video_shape[2]
+            new_height = round(video_height * (self.width / video_width))
+            new_width = self.width
+        elif self.width is None and self.height:
+            video_height = video.video_shape[1]
+            video_width = video.video_shape[2]
+            new_width = round(video_width * (self.height / video_height))
+            new_height = self.height
+
+        print(f"Resizing video to: {new_width}x{new_height}!")
         with Pool() as pool:
             frames_copy = pool.starmap(
                 self._resize_frame,
-                [(frame, self.new_width, self.new_height) for frame in video.frames],
+                [(frame, new_width, new_height) for frame in video.frames],
             )
         video.frames = np.array(frames_copy)
         return video
@@ -128,3 +147,32 @@ class ResampleFPS(Transformation):
             print(f"Upsampling video from {video.fps} to {self.new_fps} FPS.")
             video = self._upsample(video)
         return video
+
+
+class CropMode(Enum):
+    CENTER = "center"
+
+
+class Crop(Transformation):
+
+    def __init__(self, width: int, height: int, mode: CropMode = CropMode.CENTER):
+        self.width = width
+        self.height = height
+        self.mode = mode
+
+    def apply(self, video: Video) -> Video:
+        if self.mode == CropMode.CENTER:
+            current_shape = video.frame_shape[:2]
+            center_height = current_shape[0] // 2
+            center_width = current_shape[1] // 2
+            width_offset = self.width // 2
+            height_offset = self.height // 2
+            video.frames = video.frames[
+                :,
+                center_height - height_offset : center_height + height_offset,
+                center_width - width_offset : center_width + width_offset,
+                :,
+            ]
+        else:
+            raise ValueError(f"Unknown mode: {self.mode}")
+        return video

{videopython-0.1.3 → videopython-0.1.41}/src/videopython/base/transitions.py

@@ -4,6 +4,7 @@ from typing import final
 
 import numpy as np
 
+from videopython.base.effects import Blur
 from videopython.base.video import Video
 
 
@@ -68,3 +69,38 @@ class FadeTransition(Transition):
         )
         faded_videos.audio = videos[0].audio.append(videos[1].audio, crossfade=(effect_time_fps / video_fps) * 1000)
         return faded_videos
+
+
+class BlurTransition(Transition):
+    def __init__(
+        self, effect_time_seconds: float = 1.5, blur_iterations: int = 400, blur_kernel_size: tuple[int, int] = (11, 11)
+    ):
+        self.effect_time_seconds = effect_time_seconds
+        self.blur_iterations = blur_iterations
+        self.blur_kernel_size = blur_kernel_size
+
+    def _apply(self, videos: tuple[Video, Video]) -> Video:
+        video_fps = videos[0].fps
+        for video in videos:
+            if video.total_seconds < self.effect_time_seconds:
+                raise RuntimeError("Not enough space to make transition!")
+
+        effect_time_fps = math.floor(self.effect_time_seconds * video_fps)
+
+        ascending_blur = Blur("ascending", self.blur_iterations, self.blur_kernel_size)
+        descending_blur = Blur("descending", self.blur_iterations, self.blur_kernel_size)
+        transition = ascending_blur.apply(videos[0][-effect_time_fps:]) + descending_blur.apply(
+            videos[1][:effect_time_fps]
+        )
+
+        blurred_videos = Video.from_frames(
+            np.r_[
+                "0,2",
+                videos[0].frames[:-effect_time_fps],
+                transition.frames,
+                videos[1].frames[effect_time_fps:],
+            ],
+            fps=video_fps,
+        )
+        blurred_videos.audio = videos[0].audio.append(videos[1].audio)
+        return blurred_videos

{videopython-0.1.3 → videopython-0.1.41}/src/videopython/base/video.py

@@ -23,7 +23,7 @@ class VideoMetadata:
     total_seconds: float
 
     def __str__(self):
-        return f"{self.height}x{self.width} @ {self.fps}fps, {self.total_seconds} seconds"
+        return f"{self.width}x{self.height} @ {self.fps}fps, {self.total_seconds} seconds"
 
     def __repr__(self) -> str:
         return self.__str__()
@@ -122,6 +122,12 @@ class Video:
     @classmethod
     def from_frames(cls, frames: np.ndarray, fps: float) -> Video:
         new_vid = cls()
+        if frames.ndim != 4:
+            raise ValueError(f"Unsupported number of dimensions: {frames.shape}!")
+        elif frames.shape[-1] == 4:
+            frames = frames[:, :, :, :3]
+        elif frames.shape[-1] != 3:
+            raise ValueError(f"Unsupported number of dimensions: {frames.shape}!")
         new_vid.frames = frames
         new_vid.fps = fps
         new_vid.audio = AudioSegment.silent(duration=round(new_vid.total_seconds * 1000))

videopython-0.1.41/src/videopython/generation/audio.py

@@ -0,0 +1,22 @@
+import numpy as np
+import torch
+from pydub import AudioSegment
+from transformers import AutoTokenizer, VitsModel
+
+TEXT_TO_SPEECH_MODEL = "facebook/mms-tts-eng"
+
+
+class TextToSpeech:
+    def __init__(self):
+        self.pipeline = VitsModel.from_pretrained(TEXT_TO_SPEECH_MODEL)
+        self.tokenizer = AutoTokenizer.from_pretrained(TEXT_TO_SPEECH_MODEL)
+
+    def generate_audio(self, text: str) -> AudioSegment:
+        tokenized = self.tokenizer(text, return_tensors="pt")
+
+        with torch.no_grad():
+            output = self.pipeline(**tokenized).waveform
+
+        output = (output.T.float().numpy() * (2**31 - 1)).astype(np.int32)
+        audio = AudioSegment(data=output, frame_rate=self.pipeline.config.sampling_rate, sample_width=4, channels=1)
+        return audio

videopython-0.1.41/src/videopython/generation/image.py

@@ -0,0 +1,22 @@
+import io
+import os
+
+import torch
+from diffusers import DiffusionPipeline
+from PIL import Image
+
+TEXT_TO_IMAGE_MODEL = "stabilityai/stable-diffusion-xl-base-1.0"
+
+
+class TextToImage:
+    def __init__(self):
+        if not torch.cuda.is_available():
+            raise ValueError("CUDA is not available, but TextToVideo model requires CUDA.")
+        self.pipeline = DiffusionPipeline.from_pretrained(
+            TEXT_TO_IMAGE_MODEL, torch_dtype=torch.float16, variant="fp16", use_safetensors=True
+        )
+        self.pipeline.to("cuda")
+
+    def generate_image(self, prompt: str) -> Image.Image:
+        image = self.pipeline(prompt=prompt).images[0]
+        return image

{videopython-0.1.3 → videopython-0.1.41}/src/videopython/generation/video.py

@@ -10,13 +10,12 @@ IMAGE_TO_VIDEO_MODEL = "stabilityai/stable-video-diffusion-img2vid-xt"
 
 
 class TextToVideo:
-    def __init__(self, gpu_optimized: bool = True):
-        self.pipeline = DiffusionPipeline.from_pretrained(
-            TEXT_TO_VIDEO_MODEL, torch_dtype=torch.float16 if gpu_optimized else torch.float32
-        )
+    def __init__(self):
+        if not torch.cuda.is_available():
+            raise ValueError("CUDA is not available, but TextToVideo model requires CUDA.")
+        self.pipeline = DiffusionPipeline.from_pretrained(TEXT_TO_VIDEO_MODEL, torch_dtype=torch.float16)
         self.pipeline.scheduler = DPMSolverMultistepScheduler.from_config(self.pipeline.scheduler.config)
-        if gpu_optimized:
-            self.pipeline.enable_model_cpu_offload()
+        self.pipeline.to("cuda")
 
     def generate_video(
         self, prompt: str, num_steps: int = 25, height: int = 320, width: int = 576, num_frames: int = 24
@@ -39,7 +38,6 @@ class ImageToVideo:
         self.pipeline = DiffusionPipeline.from_pretrained(
             IMAGE_TO_VIDEO_MODEL, torch_dtype=torch.float16, variant="fp16"
         ).to("cuda")
-        self.pipeline.enable_model_cpu_offload()
 
     def generate_video(self, image: Image, fps: int = 24) -> Video:
         video_frames = self.pipeline(image=image, fps=fps, output_type="np").frames[0]