314 lines
14 KiB
Python
314 lines
14 KiB
Python
import torch
|
|
import torchaudio
|
|
import numpy as np
|
|
import cv2
|
|
import json
|
|
import os
|
|
|
|
|
|
|
|
def convert_to_uint8(image):
|
|
if image.dtype == np.uint8:
|
|
return image
|
|
if image.dtype != np.float32 and image.dtype != np.float64:
|
|
image = image.astype(np.float32)
|
|
if image.max() <= 1.0:
|
|
image = image * 255.0
|
|
image = np.clip(image, 0, 255)
|
|
return image.astype(np.uint8)
|
|
|
|
def convert_to_unit8_v2(image):
|
|
image = (image - image.min()) / (image.max() - image.min())*255
|
|
|
|
return image.astype(np.uint8)
|
|
|
|
|
|
|
|
|
|
def resize_image(imageTensor, target_height, target_width):
|
|
|
|
images = []
|
|
for image in imageTensor:
|
|
print("Resizing Image...")
|
|
image = image.squeeze(0).cpu().numpy()
|
|
image = convert_to_uint8(image)
|
|
image = cv2.resize(image, (int(target_width), int(target_height)), interpolation=cv2.INTER_AREA)
|
|
image = torch.from_numpy(image).permute(2, 0, 1).unsqueeze(0).float()/255
|
|
image = image.permute(0, 2, 3, 1)
|
|
images.append(image)
|
|
|
|
print ("Images resized!")
|
|
return torch.cat(images, dim=0)
|
|
|
|
|
|
|
|
|
|
def generate_silent_audio(duration, sample_rate, num_channels=2, batch_size=1):
|
|
"""
|
|
Generates a silent audio tensor of the given duration.
|
|
|
|
Parameters:
|
|
duration (float): Duration of the silence in seconds (can be fractional).
|
|
sample_rate (int): The audio sample rate (e.g., 44100 Hz).
|
|
num_channels (int): Number of audio channels (1 for mono, 2 for stereo, etc.).
|
|
batch_size (int): Number of audio batches to generate (default is 1).
|
|
|
|
Returns:
|
|
torch.Tensor: A 3D tensor filled with zeros, representing silence with shape (batch_size, num_channels, num_samples).
|
|
"""
|
|
# Compute the total number of samples based on duration and sample rate
|
|
num_samples = int(round(duration * sample_rate))
|
|
|
|
# Return a 3D tensor of zeros (representing silence)
|
|
return torch.zeros((batch_size, 2, num_samples), dtype=torch.float32)
|
|
|
|
|
|
|
|
def resample_audio(audio_tensor, target_sample_rate):
|
|
"""
|
|
Resamples the audio to the target sample rate if necessary, and ensures a 3D tensor output.
|
|
|
|
Parameters:
|
|
audio_tensor (dict): Dictionary containing 'waveform' (audio tensor) and 'sample_rate'.
|
|
target_sample_rate (int): The target sample rate to resample the audio.
|
|
|
|
Returns:
|
|
dict: Dictionary containing the resampled 'waveform' (3D tensor) and 'sample_rate'.
|
|
"""
|
|
orig_sample_rate = audio_tensor['sample_rate']
|
|
waveform = audio_tensor['waveform']
|
|
|
|
# Ensure the waveform has a batch dimension (convert 2D to 3D if necessary)
|
|
if waveform.dim() == 2: # Shape is (num_channels, num_samples)
|
|
waveform = waveform.unsqueeze(0) # Add a batch dimension -> (1, num_channels, num_samples)
|
|
|
|
# Resample if necessary
|
|
if orig_sample_rate != target_sample_rate:
|
|
resampler = torchaudio.transforms.Resample(orig_freq=orig_sample_rate, new_freq=target_sample_rate)
|
|
resampled_waveform = resampler(waveform)
|
|
return {'waveform': resampled_waveform, 'sample_rate': target_sample_rate}
|
|
else:
|
|
return {'waveform': waveform, 'sample_rate': orig_sample_rate}
|
|
|
|
|
|
|
|
def get_first_available_sample_rate(audios_list, default_sample_rate=44100):
|
|
"""
|
|
Get the sample rate of the first available (non-None) audio from the list.
|
|
|
|
Parameters:
|
|
audios_list (list of dict): List of dictionaries containing 'waveform' and 'sample_rate',
|
|
where None indicates no audio is present.
|
|
default_sample_rate (int): Default sample rate to return if no valid audio is found.
|
|
|
|
Returns:
|
|
int: The sample rate of the first available audio or the default sample rate if none is found.
|
|
"""
|
|
for audio_tensor in audios_list:
|
|
if audio_tensor is not None:
|
|
# Extract the sample rate from the dictionary
|
|
print (audio_tensor['sample_rate'])
|
|
return audio_tensor['sample_rate'] # Return the sample rate of the first valid audio
|
|
|
|
# Return the default sample rate if no valid audio is found
|
|
print("No valid audio sample rate could be determined from the audio list. Using default sample rate.")
|
|
|
|
return default_sample_rate
|
|
|
|
|
|
|
|
|
|
def create_frame_and_audio_lists(frames_set1=None, frames_set2=None, frames_set3=None, frames_set4=None, frames_set5=None, audio_set1=None, audio_set2=None, audio_set3=None, audio_set4=None, audio_set5=None):
|
|
|
|
frames_list = [frames_set1, frames_set2, frames_set3, frames_set4, frames_set5]
|
|
audios_list = [audio_set1, audio_set2, audio_set3, audio_set4, audio_set5]
|
|
filtered_frames_list = [frames for frames in frames_list if frames is not None]
|
|
|
|
return filtered_frames_list, audios_list
|
|
|
|
|
|
|
|
|
|
def combine_frames_and_audios(frames_list, audios_list=None, frame_rate=30):
|
|
"""
|
|
Combines multiple sets of image frames (as tensors) and audios (as tensors) into a single set.
|
|
Inserts silence for image sets that have no corresponding audio. Resizes and crops images.
|
|
|
|
Parameters:
|
|
frames_list (list of torch.Tensor): A list of 4D tensors of image frames (batch_size x channels x height x width).
|
|
audios_list (list of dict, optional): A list of dictionaries, each containing 'waveform' (audio tensor) and 'sample_rate', or None if no audio is present.
|
|
|
|
Returns:
|
|
combined_frames (torch.Tensor): A combined 4D tensor of image frames.
|
|
combined_audio (torch.Tensor or None): A combined 3D tensor of audio data, with silence where necessary.
|
|
"""
|
|
# Get the target width and height from the first set of frames
|
|
target_height, target_width = frames_list[0].size(1), frames_list[0].size(2)
|
|
|
|
# print(f"target_height, target_width = {frames_list[0].size(1), frames_list[0].size(2)}")
|
|
|
|
# Resize and crop each set of frames to match the first set
|
|
resized_frames_list = [resize_image(frames, target_height, target_width) for frames in frames_list]
|
|
|
|
# Combine the frames from all videos
|
|
combined_frames = torch.cat(resized_frames_list, dim=0) # Concatenate along the batch dimension
|
|
|
|
# print(video_info)
|
|
# if video_info is not None:
|
|
# first_key = next(iter(video_info))
|
|
# first_value = video_info[first_key]
|
|
|
|
fps = frame_rate
|
|
|
|
# If audios_list is provided, resample and combine audios
|
|
if audios_list is not None:
|
|
# Get the sample rate of the first available audio tensor
|
|
first_audio_sample_rate = get_first_available_sample_rate(audios_list)
|
|
|
|
# Prepare a list to store resampled or silent audio tensors
|
|
combined_audios = []
|
|
combined_audios_dict = []
|
|
|
|
for i, frames in enumerate(frames_list):
|
|
|
|
num_frames = frames.size(0) # Number of frames in this set
|
|
print(f"There are {num_frames} frames in this video.")
|
|
duration = num_frames / fps # Duration in seconds based on the frame rate
|
|
|
|
# print(f"The duration of this frame is : {duration}")
|
|
|
|
audio_dict = audios_list[i] # The audio dictionary for this set
|
|
|
|
if audio_dict is None:
|
|
# Extract the number of channels from the first available audio or default to 1
|
|
num_channels = audios_list[0]['waveform'].size(0) if audios_list[0] is not None else 2
|
|
silent_audio = generate_silent_audio(duration, first_audio_sample_rate, num_channels)
|
|
combined_audios.append(silent_audio)
|
|
|
|
print("No audio provided, so silent audio added!")
|
|
# print(f"Silent Audio: {silent_audio.shape}")
|
|
else:
|
|
# Extract the waveform and sample_rate from the audio dictionary
|
|
waveform = audio_dict['waveform']
|
|
sample_rate = audio_dict['sample_rate']
|
|
print(f"Audio provided! Sample rate: {sample_rate}")
|
|
# Resample the audio to match the sample rate of the first audio set
|
|
resampled_audio = resample_audio({'waveform': waveform, 'sample_rate': sample_rate}, first_audio_sample_rate)
|
|
print(f"New Sample rate: {resampled_audio['sample_rate']}")
|
|
# print(f"Audio: {resampled_audio['waveform'].shape}")
|
|
combined_audios.append(resampled_audio['waveform']) # Append the resampled waveform
|
|
# combined_audios.append(audio_dict['waveform'])
|
|
|
|
# Concatenate all audios (real and silent)
|
|
combined_audio = torch.cat(combined_audios, dim=2) # Concatenate along the sample dimension
|
|
|
|
combined_audios_dict = {
|
|
'waveform': combined_audio,
|
|
'sample_rate': first_audio_sample_rate
|
|
}
|
|
|
|
else:
|
|
combined_audios_dict = None
|
|
print("No audio provided!")
|
|
|
|
return combined_frames, combined_audios_dict
|
|
|
|
|
|
|
|
def get_script_dir():
|
|
"""Get the directory where the script is located."""
|
|
return os.path.dirname(os.path.abspath(__file__))
|
|
|
|
def load_styles():
|
|
"""Load the styles.json file from the script's directory."""
|
|
script_dir = get_script_dir()
|
|
styles_path = os.path.join(script_dir, "styles.json")
|
|
|
|
try:
|
|
with open(styles_path, "r") as f:
|
|
return json.load(f)
|
|
except FileNotFoundError:
|
|
print(f"Error: The styles.json file was not found in {styles_path}.")
|
|
# Optional: return default styles as a fallback
|
|
return {
|
|
"none": {
|
|
"positive": "",
|
|
"negative": ""
|
|
},
|
|
"anime1": {
|
|
"positive": "Japanese comic book anime style, samurai x, exaggerated features, sharp dark lines, anime, anime style, colorful, vibrant, detailed, intricate, large expressive eyes",
|
|
"negative": "text, logo, signature, watermark, soft edges, low quality, blurry, realistic style, distorted faces, bad anatomy, poor proportions"
|
|
},
|
|
"anime2": {
|
|
"positive": "a beautiful anime-style illustration, expressive eyes, exaggerated features, vibrant colors, detailed, expressive characters, studio ghibli influence",
|
|
"negative": "low quality, blurry, realistic style, distorted faces, bad anatomy, poor proportions"
|
|
},
|
|
"beauArt": {
|
|
"positive": "hard edges, black and white book illustration, bold lines, black lines, white background",
|
|
"negative": "color, poor contrast, realistic, soft edges"
|
|
},
|
|
"black Line art":{
|
|
"positive": "Black Line art, black and white, sharp lines, clear, minimalist, graphic, precision, ink, monochrome",
|
|
"negative": "photograph, stock photo, realistic, deformed, glitch, color, vague, blurry, noisy, low contrast, photorealistic, realism, impressionism, expressionism, oil, acrylic, watercolor, pastel, textured, gradient, shaded"
|
|
},
|
|
"comic1": {
|
|
"positive": "vivid comic book illustration, bold lines, halftone shading, dramatic action",
|
|
"negative": "washed-out colors, lack of contrast, soft edges, overly detailed backgrounds"
|
|
},
|
|
"comic2": {
|
|
"positive": "comic book style, dynamic pose, exaggerated features, bold lines, vibrant colors, ink wash, retro style, comic book panel",
|
|
"negative": "blurry, low quality, poorly drawn, deformed, disfigured, extra limbs, missing limbs, extra fingers, missing fingers, unrealistic anatomy, poorly proportioned, monochrome, grayscale, black and white"
|
|
},
|
|
"disney": {
|
|
"positive": "classic Disney-style art, colorful, detailed characters, expressive emotions, smooth animations, magical and whimsical setting",
|
|
"negative": "dull colors, lack of detail, distorted features, dark themes"
|
|
},
|
|
"manga": {
|
|
"positive": "black and white manga illustration, bold lines, dynamic compositions, detailed linework, expressive characters, action-packed panels",
|
|
"negative": "color, poor contrast, lack of detail, blurry, distorted anatomy, flat expressions"
|
|
},
|
|
"ps1": {
|
|
"positive": "low-poly 3D graphics in PS1 style, retro gaming aesthetics, simple geometric shapes, nostalgic vibe",
|
|
"negative": "modern high-resolution graphics, realistic textures, excessive detail, smooth rendering"
|
|
},
|
|
"pixar": {
|
|
"positive": "a 3D animation in Pixar style, highly detailed, vibrant, whimsical, expressive characters, cinematic lighting, storybook atmosphere",
|
|
"negative": "grainy textures, flat colors, dull, uninspired composition, poor rendering"
|
|
},
|
|
"toy": {
|
|
"positive": "toy-like miniature characters, bright colors, soft plastic textures, whimsical design, playful atmosphere, detailed craftsmanship",
|
|
"negative": "realistic details, dull colors, rough textures, lack of charm"
|
|
},
|
|
"watercolor": {
|
|
"positive": "soft and vibrant watercolor painting, gentle gradients, flowing textures, light and airy feel, pastel tones",
|
|
"negative": "harsh lines, digital appearance, over-saturated colors, poor blending"
|
|
}
|
|
|
|
}
|
|
|
|
|
|
def get_available_styles():
|
|
styles = load_styles()
|
|
return list(styles.keys())
|
|
|
|
|
|
def generate_prompts(art_style, user_positive_prompt="sharp image", user_negative_prompt="blurry"):
|
|
styles = load_styles()
|
|
|
|
if art_style not in styles:
|
|
return {
|
|
"positive_prompt": f"Sorry, no predefined prompts for the '{art_style}' style. Please specify another style.",
|
|
"negative_prompt": "None"
|
|
}
|
|
|
|
style = styles[art_style]
|
|
combined_positive = f"{style['positive']}, {user_positive_prompt}".strip(", ")
|
|
combined_negative = f"{style['negative']}, {user_negative_prompt}".strip(", ")
|
|
|
|
return {
|
|
"positive_prompt": combined_positive,
|
|
"negative_prompt": combined_negative
|
|
}
|
|
|
|
|