106 lines
4.0 KiB
Python
106 lines
4.0 KiB
Python
import os
|
|
import torch
|
|
from diffusers import DiffusionPipeline, DPMSolverMultistepScheduler
|
|
from diffusers.utils import export_to_video
|
|
from PIL import Image, ImageOps
|
|
import numpy as np
|
|
|
|
class ImageDirIterator:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"directory_path": ("STRING", {}),
|
|
"image_index": ("INT", {"default": 0})
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "get_image_by_index"
|
|
CATEGORY = "cspnodes"
|
|
|
|
def get_image_by_index(self, directory_path, image_index):
|
|
# Get list of image files sorted by modification time (most recent first)
|
|
image_files = sorted(
|
|
[os.path.join(directory_path, f) for f in os.listdir(directory_path) if f.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif'))],
|
|
key=lambda x: os.path.getmtime(x),
|
|
reverse=True
|
|
)
|
|
|
|
# Validate index
|
|
if image_index < 0 or image_index >= len(image_files):
|
|
raise IndexError("Image index out of range.")
|
|
|
|
# Load and preprocess the image
|
|
image = Image.open(image_files[image_index])
|
|
image = ImageOps.exif_transpose(image) # Correct orientation
|
|
image = image.convert("RGB") # Ensure image is in RGB format
|
|
|
|
# Convert image to tensor
|
|
image_tensor = torch.from_numpy(np.array(image).astype(np.float32) / 255.0)[None,]
|
|
|
|
return (image_tensor,)
|
|
|
|
class Modelscopet2v:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"prompt": ("STRING", {}),
|
|
"negative_prompt": ("STRING", {"default": None}),
|
|
"num_inference_steps": ("INT", {"default": 40}),
|
|
"guidance_scale": ("FLOAT", {"default": 7.5}),
|
|
"seed": ("INT", {"default": None}),
|
|
"height": ("INT", {"default": 320}),
|
|
"width": ("INT", {"default": 576}),
|
|
"num_frames": ("INT", {"default": 24}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "generate_video_frames"
|
|
CATEGORY = "cspnodes"
|
|
|
|
def generate_video_frames(self, prompt, num_inference_steps, height, width, num_frames, guidance_scale, negative_prompt, seed):
|
|
# Set up the generator for deterministic results if seed is provided
|
|
generator = torch.Generator()
|
|
if seed is not None:
|
|
generator.manual_seed(seed)
|
|
|
|
pipe = DiffusionPipeline.from_pretrained("cerspense/zeroscope_v2_576w", torch_dtype=torch.float16)
|
|
pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
|
|
pipe.enable_model_cpu_offload()
|
|
|
|
# Added generator to the pipe call
|
|
video_frames = pipe(prompt, num_inference_steps=num_inference_steps, height=height, width=width, num_frames=num_frames, guidance_scale=guidance_scale, negative_prompt=negative_prompt, generator=generator).frames
|
|
|
|
# Print the shape of the video frames to debug
|
|
print(f"Shape of the video frames: {video_frames.shape}")
|
|
|
|
# Ensure video_frames is a PyTorch tensor
|
|
if not isinstance(video_frames, torch.Tensor):
|
|
video_frames = torch.tensor(video_frames, dtype=torch.float32)
|
|
|
|
# Remove the unnecessary batch dimension explicitly and permute the dimensions
|
|
# The expected shape is (num_frames, height, width, channels)
|
|
video_frames = video_frames.squeeze(0).permute(0, 1, 2, 3)
|
|
|
|
# Convert the tensor to CPU and to uint8 if it's not already
|
|
video_frames = video_frames.to('cpu')
|
|
|
|
# Print the shape of the video frames tensor to debug
|
|
print(f"Shape of the video frames tensor: {video_frames.shape}")
|
|
|
|
# return (video_frames_numpy,)
|
|
return (video_frames,)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"ImageDirIterator": ImageDirIterator,
|
|
"Modelscopet2v": Modelscopet2v # Added new node class
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"ImageDirIterator": "Image Dir Iterator",
|
|
"Modelscopet2v": "Modelscopet2v" # Added display name for new node
|
|
} |