diff --git a/runner/app/test_examples/continuous-creation.py b/runner/app/test_examples/continuous-creation.py new file mode 100644 index 00000000..54473930 --- /dev/null +++ b/runner/app/test_examples/continuous-creation.py @@ -0,0 +1,154 @@ +import torch +import numpy as np +import os +import shutil +import time +import sys +sys.path.append('C://Users//ganes//ai-worker//runner') +from PIL import PngImagePlugin, Image +from diffusers import StableVideoDiffusionPipeline +from app.pipelines.frame_interpolation import FILMPipeline +from app.pipelines.utils import DirectoryReader, frames_compactor, DirectoryWriter + +# Increase the max text chunk size for PNG images +PngImagePlugin.MAX_TEXT_CHUNK = 100 * (1024**2) + +def move_file_with_retry(src_file_path: str, dst_file_path: str, retries: int = 5, delay: float = 1.0): + for attempt in range(retries): + try: + shutil.move(src_file_path, dst_file_path) + return + except PermissionError: + print(f"Attempt {attempt + 1} failed: File is in use. Retrying in {delay} seconds...") + time.sleep(delay) + raise PermissionError(f"Failed to move file after {retries} attempts.") + +def get_last_file_sorted_by_name(directory: str) -> str: + """ + Get the last file in the directory when sorted by filename. + + Args: + directory (str): Path to the directory. + + Returns: + str: Path to the last file in the sorted list, or None if directory is empty. + """ + try: + # List all files in the directory + files = [f for f in os.listdir(directory) if os.path.isfile(os.path.join(directory, f))] + + if not files: + print("No files found in the directory.") + return None + + # Sort files by name + files.sort() + + # Get the last file in the sorted list + last_file = files[-1] + + return os.path.join(directory, last_file) + + except Exception as e: + print(f"An error occurred: {e}") + return None + +def main(): + # Initialize pipelines + repo_id = "stabilityai/stable-video-diffusion-img2vid-xt" + svd_xt_pipeline = StableVideoDiffusionPipeline.from_pretrained( + repo_id, torch_dtype=torch.float16, variant="fp16", cache_dir="G:/ai-models/models" + ) + svd_xt_pipeline.enable_model_cpu_offload() + + film_pipeline = FILMPipeline("G:/ai-models/models/film_net_fp16.pt").to(device="cuda", dtype=torch.float16) + + # Load initial input image + image_path = "G:/ai-models/models/gif_frames/donut_motion.png" + image = Image.open(image_path) + + fps = 24.0 + inter_frames = 4 + rounds = 2 # Number of rounds of generation and interpolation + base_output_dir = "G:/ai-models/models" + + all_frames_dir = os.path.join(base_output_dir, "all_interpolated_frames") + os.makedirs(all_frames_dir, exist_ok=True) + + last_frame_for_next_round = os.path.join(base_output_dir, "last_frame_for_next_round.png") + + for round_num in range(1, rounds + 1): + svd_xt_output_dir = os.path.join(base_output_dir, f"svd_xt_output_round_{round_num}") + os.makedirs(svd_xt_output_dir, exist_ok=True) + + # Generate frames using SVD pipeline + generator = torch.manual_seed(42) + frames = svd_xt_pipeline(image, decode_chunk_size=8, generator=generator, output_type="np").frames[0] + + # Save SVD frames to directory + film_writer = DirectoryWriter(svd_xt_output_dir) + for idx, frame in enumerate(frames): + film_writer.write_frame(torch.tensor(frame).permute(2, 0, 1)) + + # Read saved frames for interpolation + film_reader = DirectoryReader(svd_xt_output_dir) + height, width = film_reader.get_resolution() + + # Interpolate frames using FILM pipeline + film_pipeline(film_reader, film_writer, inter_frames=inter_frames) + + # Close reader and writer + film_writer.close() + film_reader.reset() + + # Deleting the SVD generated images. + for i in range(len(frames)): + os.remove(os.path.join(svd_xt_output_dir, f"{i}.png")) + print(f"Deleted the first {len(frames)} frames in directory: {svd_xt_output_dir}") + + # Save the last frame separately for the next round + last_frame_path = get_last_file_sorted_by_name(svd_xt_output_dir) + if last_frame_path: + shutil.copy2(last_frame_path, last_frame_for_next_round) + else: + print("No frames found to copy.") + + # Move all interpolated frames to a common directory with a unique naming scheme + for file_name in sorted(os.listdir(svd_xt_output_dir)): + src_file_path = os.path.join(svd_xt_output_dir, file_name) + dst_file_name = f"round_{round_num:03d}_frame_{file_name}" + dst_file_path = os.path.join(all_frames_dir, dst_file_name) + + move_file_with_retry(src_file_path, dst_file_path) + + # Clean up the source directory after moving frames + for file_name in os.listdir(svd_xt_output_dir): + os.remove(os.path.join(svd_xt_output_dir, file_name)) + os.rmdir(svd_xt_output_dir) + + # Ensure all operations on last frame are complete before opening it again + time.sleep(1) # Small delay to ensure file system operations are complete + + # Prepare for next round + image = Image.open(last_frame_for_next_round) + + # Compile all interpolated frames from all rounds into a final video + video_output_dir = "G:/ai-models/models/video_out" + os.makedirs(video_output_dir, exist_ok=True) + + film_output_path = os.path.join(video_output_dir, "output.avi") + frames_compactor(frames=all_frames_dir, output_path=film_output_path, fps=fps, codec="MJPG", is_directory=True) + + # Clean up all frames in the directories after video generation + for file_name in os.listdir(all_frames_dir): + file_path = os.path.join(all_frames_dir, file_name) + if os.path.isfile(file_path): + os.remove(file_path) + os.rmdir(all_frames_dir) + + print(f"All frames deleted from directories.") + print(f"Video generated at: {film_output_path}") + return film_output_path + +if __name__ == "__main__": + main() diff --git a/runner/app/test_examples/generate-interpolate-upscale.py b/runner/app/test_examples/generate-interpolate-upscale.py new file mode 100644 index 00000000..69954a0a --- /dev/null +++ b/runner/app/test_examples/generate-interpolate-upscale.py @@ -0,0 +1,159 @@ +import logging +import os +os.environ["MODEL_DIR"] = "G://ai-models//models" +import shutil +import time +import sys +sys.path.append('C://Users//ganes//ai-worker//runner') +from typing import List, Optional, Tuple +import PIL +import torch +from diffusers import StableVideoDiffusionPipeline +from PIL import PngImagePlugin, Image, ImageFile + +from app.pipelines.base import Pipeline +from app.pipelines.frame_interpolation import FILMPipeline +from app.pipelines.upscale import UpscalePipeline +from app.pipelines.utils import DirectoryReader, frames_compactor, DirectoryWriter, SafetyChecker, get_model_dir, get_torch_device, is_lightning_model, is_turbo_model +from huggingface_hub import file_download + +# Increase the max text chunk size for PNG images +PngImagePlugin.MAX_TEXT_CHUNK = 100 * (1024**2) +ImageFile.LOAD_TRUNCATED_IMAGES = True + +logger = logging.getLogger(__name__) + +# Helper function to move files with retry mechanism +def move_file_with_retry(src_file_path: str, dst_file_path: str, retries: int = 5, delay: float = 1.0): + for attempt in range(retries): + try: + shutil.move(src_file_path, dst_file_path) + return + except PermissionError: + print(f"Attempt {attempt + 1} failed: File is in use. Retrying in {delay} seconds...") + time.sleep(delay) + raise PermissionError(f"Failed to move file after {retries} attempts.") + +# Helper function to get the last file in a directory sorted by filename +def get_last_file_sorted_by_name(directory: str) -> str: + try: + files = [f for f in os.listdir(directory) if os.path.isfile(os.path.join(directory, f))] + if not files: + print("No files found in the directory.") + return None + files.sort() + last_file = files[-1] + return os.path.join(directory, last_file) + except Exception as e: + print(f"An error occurred: {e}") + return None + +def main(): + # Initialize SVD and FILM pipelines + repo_id = "stabilityai/stable-video-diffusion-img2vid-xt" + svd_xt_pipeline = StableVideoDiffusionPipeline.from_pretrained( + repo_id, torch_dtype=torch.float16, variant="fp16", cache_dir="G:/ai-models/models" + ) + svd_xt_pipeline.enable_model_cpu_offload() + + film_pipeline = FILMPipeline("G:/ai-models/models/film_net_fp16.pt").to(device="cuda", dtype=torch.float16) + + # Load initial input image + image_path = "G:/ai-models/models/gif_frames/donut_motion.png" + image = Image.open(image_path) + + fps = 24.0 + inter_frames = 4 + rounds = 2 # Number of rounds of generation and interpolation + base_output_dir = "G:/ai-models/models" + + all_frames_dir = os.path.join(base_output_dir, "all_interpolated_frames") + os.makedirs(all_frames_dir, exist_ok=True) + + last_frame_for_next_round = os.path.join(base_output_dir, "last_frame_for_next_round.png") + + for round_num in range(1, rounds + 1): + svd_xt_output_dir = os.path.join(base_output_dir, f"svd_xt_output_round_{round_num}") + os.makedirs(svd_xt_output_dir, exist_ok=True) + + # Generate frames using SVD pipeline + generator = torch.manual_seed(42) + frames = svd_xt_pipeline(image, decode_chunk_size=8, generator=generator, output_type="np").frames[0] + + # Save SVD frames to directory + film_writer = DirectoryWriter(svd_xt_output_dir) + for idx, frame in enumerate(frames): + film_writer.write_frame(torch.tensor(frame).permute(2, 0, 1)) + + # Read saved frames for interpolation + film_reader = DirectoryReader(svd_xt_output_dir) + height, width = film_reader.get_resolution() + + # Interpolate frames using FILM pipeline + film_pipeline(film_reader, film_writer, inter_frames=inter_frames) + + # Close reader and writer + film_writer.close() + film_reader.reset() + + # Deleting the SVD generated images. + for i in range(len(frames)): + os.remove(os.path.join(svd_xt_output_dir, f"{i}.png")) + print(f"Deleted the first {len(frames)} frames in directory: {svd_xt_output_dir}") + + # Save the last frame separately for the next round + last_frame_path = get_last_file_sorted_by_name(svd_xt_output_dir) + if last_frame_path: + shutil.copy2(last_frame_path, last_frame_for_next_round) + else: + print("No frames found to copy.") + + # Initialize Upscale pipeline and Upscale the last frame before passing to the next round + upscale_pipeline = UpscalePipeline("stabilityai/stable-diffusion-x4-upscaler", torch_dtype=torch.float16) + upscale_pipeline.enable_model_cpu_offload() + upscale_pipeline.sfast_enabled() + upscaled_image, _ = upscale_pipeline("", image=Image.open(last_frame_for_next_round),) + print('Upscaling of the seed image before next round.') + print(upscaled_image[0].shape) + exit + upscaled_image[0].save(last_frame_for_next_round) + + # Move all interpolated frames to a common directory with a unique naming scheme + for file_name in sorted(os.listdir(svd_xt_output_dir)): + src_file_path = os.path.join(svd_xt_output_dir, file_name) + dst_file_name = f"round_{round_num:03d}_frame_{file_name}" + dst_file_path = os.path.join(all_frames_dir, dst_file_name) + + move_file_with_retry(src_file_path, dst_file_path) + + # Clean up the source directory after moving frames + for file_name in os.listdir(svd_xt_output_dir): + os.remove(os.path.join(svd_xt_output_dir, file_name)) + os.rmdir(svd_xt_output_dir) + + # Ensure all operations on last frame are complete before opening it again + time.sleep(1) # Small delay to ensure file system operations are complete + + # Prepare for next round + image = Image.open(last_frame_for_next_round) + + # Compile all interpolated frames from all rounds into a final video + video_output_dir = "G:/ai-models/models/video_out" + os.makedirs(video_output_dir, exist_ok=True) + + film_output_path = os.path.join(video_output_dir, "output.avi") + frames_compactor(frames=all_frames_dir, output_path=film_output_path, fps=fps, codec="MJPG", is_directory=True) + + # Clean up all frames in the directories after video generation + for file_name in os.listdir(all_frames_dir): + file_path = os.path.join(all_frames_dir, file_name) + if os.path.isfile(file_path): + os.remove(file_path) + os.rmdir(all_frames_dir) + + print(f"All frames deleted from directories.") + print(f"Video generated at: {film_output_path}") + return film_output_path + +if __name__ == "__main__": + main() diff --git a/runner/app/test_examples/test-film-interpolate.py b/runner/app/test_examples/test-film-interpolate.py new file mode 100644 index 00000000..195a3a4b --- /dev/null +++ b/runner/app/test_examples/test-film-interpolate.py @@ -0,0 +1,45 @@ +import os +import requests + +# Define the URL of the FastAPI application +URL = "http://localhost:8000/FILMPipeline" + +# Test with two images +def test_with_two_images(): + image1_path = "G:/ai-models/models/all_interpolated_frames/round_003_frame_74.png" + image2_path = "G:/ai-models/models/all_interpolated_frames/round_003_frame_36.png" + + with open(image1_path, "rb") as image1, open(image2_path, "rb") as image2: + files = { + "image1": ("image1.png", image1, "image/png"), + "image2": ("image2.png", image2, "image/png"), + } + data = { + "inter_frames": 2, + "model_id": "film_net_fp16.pt" + } + response = requests.post(URL, files=files, data=data) + + print("Test with two images") + print(response.status_code) + print(response.json()) + +# Test with a directory of images +def test_with_image_directory(): + image_dir = "path/to/image_directory" + + data = { + "inter_frames": 2, + "model_path": "path/to/film_net_fp16.pt", + "image_dir": image_dir + } + response = requests.post(URL, data=data) + + print("Test with image directory") + print(response.status_code) + print(response.json()) + +if __name__ == "__main__": + # Ensure that the FastAPI server is running before executing these tests + test_with_two_images() + test_with_image_directory() diff --git a/runner/app/test_examples/test-optimized.py b/runner/app/test_examples/test-optimized.py new file mode 100644 index 00000000..db200286 --- /dev/null +++ b/runner/app/test_examples/test-optimized.py @@ -0,0 +1,69 @@ +import torch +import numpy as np +import os +import sys +sys.path.append('C://Users//ganes//ai-worker//runner') +from diffusers import StableVideoDiffusionPipeline +from PIL import PngImagePlugin, Image +from app.pipelines.frame_interpolation import FILMPipeline +from app.pipelines.utils import DirectoryReader, frames_compactor, DirectoryWriter + +# Increase the max text chunk size for PNG images +PngImagePlugin.MAX_TEXT_CHUNK = 100 * (1024**2) + +def main(): + # Initialize pipelines + repo_id = "stabilityai/stable-video-diffusion-img2vid-xt" + svd_xt_pipeline = StableVideoDiffusionPipeline.from_pretrained( + repo_id, torch_dtype=torch.float16, variant="fp16", cache_dir="G:/ai-models/models" + ) + svd_xt_pipeline.enable_model_cpu_offload() + + film_pipeline = FILMPipeline("G:/ai-models/models/film_net_fp16.pt").to(device="cuda", dtype=torch.float16) + + # Load input image + image_path = "G:/ai-models/models/gif_frames/rocket.png" + image = Image.open(image_path) + + # Generate frames using SVD pipeline + generator = torch.manual_seed(42) + frames = svd_xt_pipeline(image, decode_chunk_size=8, generator=generator, output_type="np").frames[0] + + fps = 24.0 + inter_frames = 2 + svd_xt_output_dir = "G:/ai-models/models/svd_xt_output" + video_output_dir = "G:/ai-models/models/video_out" + + # Save SVD frames to directory + film_writer = DirectoryWriter(svd_xt_output_dir) + for frame in frames: + film_writer.write_frame(torch.tensor(frame).permute(2, 0, 1)) + + # Read saved frames for interpolation + film_reader = DirectoryReader(svd_xt_output_dir) + height, width = film_reader.get_resolution() + + # Interpolate frames using FILM pipeline + film_pipeline(film_reader, film_writer, inter_frames=inter_frames) + + # Delete original SVD frames since interpolated frames are also in the same directory. + for i in range(len(frames)): + os.remove(os.path.join(svd_xt_output_dir, f"{i}.png")) + print(f"Deleted the first {len(frames)} frames in directory: {svd_xt_output_dir}") + + # Compile interpolated frames into a video + film_output_path = os.path.join(video_output_dir, "output.avi") + frames_compactor(frames=svd_xt_output_dir, output_path=film_output_path, fps=fps, codec="MJPG", is_directory=True) + + # Clean up all frames in the directory after video generation + for file_name in os.listdir(svd_xt_output_dir): + file_path = os.path.join(svd_xt_output_dir, file_name) + if os.path.isfile(file_path): + os.remove(file_path) + print(f"All frames deleted from directory: {svd_xt_output_dir}") + + print(f"Video generated at: {film_output_path}") + return film_output_path + +if __name__ == "__main__": + main() diff --git a/runner/app/test_examples/utility-func-compactor.py b/runner/app/test_examples/utility-func-compactor.py new file mode 100644 index 00000000..6dac4309 --- /dev/null +++ b/runner/app/test_examples/utility-func-compactor.py @@ -0,0 +1,89 @@ +import subprocess +import numpy as np +import torch +import os +from typing import List, Union +from pathlib import Path +import cv2 + +def generate_video_from_frames( + frames: Union[List[np.ndarray], List[torch.Tensor]], + output_path: str, + fps: float, + codec: str = "MJPEG", + is_directory: bool = False, + width: int = None, + height: int = None +) -> None: + """ + Generate a video from a list of frames. Frames can be from a directory or in-memory. + + Args: + frames (List[np.ndarray] | List[torch.Tensor]): List of frames as NumPy arrays or PyTorch tensors. + output_path (str): Path to save the output video file. + fps (float): Frames per second for the video. + codec (str): Codec used for video compression (default is "MJPEG"). + is_directory (bool): If True, treat `frames` as a directory path containing image files. + width (int): Width of the video. Must be provided if `frames` are in-memory. + height (int): Height of the video. Must be provided if `frames` are in-memory. + + Returns: + None + """ + if is_directory: + # Read frames from a directory + frames = [cv2.imread(os.path.join(frames, file)) for file in sorted(os.listdir(frames))] + else: + # Convert torch tensors to numpy arrays if necessary + if isinstance(frames[0], torch.Tensor): + frames = [frame.permute(1, 2, 0).cpu().numpy() for frame in frames] + + # Ensure frames are numpy arrays and are uint8 type + frames = [frame.astype(np.uint8) for frame in frames] + + # Check if frames are consistent + if not frames: + raise ValueError("No frames to process.") + + if width is None or height is None: + # Use dimensions of the first frame if not provided + height, width = frames[0].shape[:2] + + # Write frames to a temporary directory + temp_dir = Path("temp_frames") + temp_dir.mkdir(exist_ok=True) + for i, frame in enumerate(frames): + cv2.imwrite(str(temp_dir / f"frame_{i:05d}.png"), frame) + + # Build ffmpeg command + ffmpeg_cmd = [ + 'ffmpeg', '-y', '-framerate', str(fps), + '-i', str(temp_dir / 'frame_%05d.png'), + '-c:v', codec, '-pix_fmt', 'yuv420p', + output_path + ] + + # Run ffmpeg command + subprocess.run(ffmpeg_cmd, check=True) + + # Clean up temporary frames + for file in temp_dir.glob("*.png"): + file.unlink() + temp_dir.rmdir() + + print(f"Video saved to {output_path}") + +# Example usage +if __name__ == "__main__": + # Example with in-memory frames (as np.ndarray) + # Assuming `in_memory_frames` is a list of numpy arrays + + # Example with frames from a directory + frames_directory = "G:/ai-models/models/svd_xt_output" + generate_video_from_frames( + frames=frames_directory, + output_path="G:/ai-models/models/video_out/output.mp4", + fps=24.0, + codec="mpeg4", + is_directory=True + ) diff --git a/runner/app/test_examples/utility-function-vid.py b/runner/app/test_examples/utility-function-vid.py new file mode 100644 index 00000000..adf194cc --- /dev/null +++ b/runner/app/test_examples/utility-function-vid.py @@ -0,0 +1,62 @@ +import cv2 +import numpy as np +import tempfile +import os +from io import BytesIO + +def extract_frames_from_video(video_data, is_file_path=True) -> np.ndarray: + """ + Extract frames from a video file or in-memory video data and return them as a NumPy array. + + Args: + video_data (str or BytesIO): Path to the input video file or in-memory video data. + is_file_path (bool): Indicates if video_data is a file path (True) or in-memory data (False). + + Returns: + np.ndarray: Array of frames with shape (num_frames, height, width, channels). + """ + if is_file_path: + # Handle file-based video input + video_capture = cv2.VideoCapture(video_data) + else: + # Handle in-memory video input + # Create a temporary file to store in-memory video data + with tempfile.NamedTemporaryFile(delete=False, suffix='.mp4') as temp_file: + temp_file.write(video_data.getvalue()) + temp_file_path = temp_file.name + + # Open the temporary video file + video_capture = cv2.VideoCapture(temp_file_path) + + if not video_capture.isOpened(): + raise ValueError("Error opening video data") + + frames = [] + success, frame = video_capture.read() + + while success: + frames.append(frame) + success, frame = video_capture.read() + + video_capture.release() + + # Delete the temporary file if it was created + if not is_file_path: + os.remove(temp_file_path) + + # Convert list of frames to a NumPy array + frames_array = np.array(frames) + print(f"Extracted {frames_array.shape[0]} frames from video in shape of {frames_array.shape}") + + return frames_array + +# Example usage +if __name__ == "__main__": + # File path example + video_file_path = "C:/Users/ganes/Desktop/Generated videos/output.mp4" + + + # In-memory video example + with open(video_file_path, "rb") as f: + video_data = BytesIO(f.read()) + frames_array_from_memory = extract_frames_from_video(video_data, is_file_path=False) diff --git a/runner/uvicorn.env b/runner/uvicorn.env new file mode 100644 index 00000000..77869fee --- /dev/null +++ b/runner/uvicorn.env @@ -0,0 +1,4 @@ +# myenvfile.env +MODEL_ID="" +MODEL_DIR="" +PIPELINE=FILMPipeline