355 lines
16 KiB
Python
355 lines
16 KiB
Python
# liquid_time_node.py
|
|
|
|
import os
|
|
import torch
|
|
import numpy as np
|
|
import urllib.request
|
|
from tqdm import tqdm
|
|
import math
|
|
|
|
class DownloadProgressBar(tqdm):
|
|
def update_to(self, b=1, bsize=1, tsize=None):
|
|
if tsize is not None:
|
|
self.total = tsize
|
|
self.update(b * bsize - self.n)
|
|
|
|
class LiquidTimeNode:
|
|
def __init__(self):
|
|
self.model = None
|
|
self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
|
self.precision = torch.float16 if torch.cuda.is_available() else torch.float32
|
|
|
|
print(f"[LiquidTime] Initializing with device: {self.device}, precision: {self.precision}")
|
|
|
|
self.model_urls = {
|
|
'fp16': 'https://github.com/dajes/frame-interpolation-pytorch/releases/download/v1.0.2/film_net_fp16.pt',
|
|
'fp32': 'https://github.com/dajes/frame-interpolation-pytorch/releases/download/v1.0.2/film_net_fp32.pt'
|
|
}
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"target_frames": ("INT", {"default": 60, "min": 2, "max": 100000}),
|
|
"start_curve": ("INT", {"default": 0, "min": -100, "max": 100}),
|
|
"middle_curve": ("INT", {"default": 0, "min": -100, "max": 100}),
|
|
"middle_moment": ("INT", {"default": 50, "min": 0, "max": 100}),
|
|
"end_curve": ("INT", {"default": 0, "min": -100, "max": 100})
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "interpolate_frames"
|
|
CATEGORY = "animation/LiquidTime"
|
|
|
|
def download_model(self, model_path):
|
|
"""Download model with progress bar"""
|
|
model_type = 'fp16' if self.precision == torch.float16 else 'fp32'
|
|
url = self.model_urls[model_type]
|
|
|
|
print(f"\n[LiquidTime] Creating model directory...")
|
|
os.makedirs(os.path.dirname(model_path), exist_ok=True)
|
|
|
|
print(f"[LiquidTime] Starting download of {model_type} model from {url}")
|
|
try:
|
|
with DownloadProgressBar(unit='B', unit_scale=True, miniters=1,
|
|
desc=f"[LiquidTime] Downloading {model_type} model",
|
|
bar_format='{desc}: {percentage:3.0f}%|{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}]') as t:
|
|
urllib.request.urlretrieve(
|
|
url,
|
|
filename=model_path,
|
|
reporthook=t.update_to
|
|
)
|
|
print(f"[LiquidTime] Successfully downloaded {model_type} model to {model_path}")
|
|
except Exception as e:
|
|
print(f"[LiquidTime] Error downloading model: {str(e)}")
|
|
if os.path.exists(model_path):
|
|
os.remove(model_path)
|
|
raise
|
|
|
|
def get_model_path(self):
|
|
"""Get path to model and ensure it exists"""
|
|
base_path = os.path.dirname(os.path.abspath(__file__))
|
|
models_dir = os.path.join(base_path, 'models')
|
|
os.makedirs(models_dir, exist_ok=True)
|
|
|
|
model_type = 'fp16' if self.precision == torch.float16 else 'fp32'
|
|
model_path = os.path.join(models_dir, f'film_net_{model_type}.pt')
|
|
|
|
if not os.path.exists(model_path):
|
|
print(f"[LiquidTime] Model not found at {model_path}")
|
|
self.download_model(model_path)
|
|
|
|
return model_path
|
|
|
|
def load_model(self):
|
|
if self.model is None:
|
|
print("[LiquidTime] Loading model...")
|
|
model_path = self.get_model_path()
|
|
try:
|
|
self.model = torch.jit.load(model_path, map_location='cpu')
|
|
self.model.eval().to(device=self.device, dtype=self.precision)
|
|
print(f"[LiquidTime] Successfully loaded model with precision {self.precision} on {self.device}")
|
|
except Exception as e:
|
|
print(f"[LiquidTime] Error loading model: {str(e)}")
|
|
raise RuntimeError(f"Error loading model: {str(e)}")
|
|
|
|
def sample_frames(self, source_frames, target_count):
|
|
"""
|
|
Wybiera target_count klatek z source_frames w równych odstępach czasowych.
|
|
"""
|
|
if target_count >= len(source_frames):
|
|
return source_frames
|
|
|
|
print(f"[LiquidTime] Reducing frames from {len(source_frames)} to {target_count}")
|
|
|
|
# Obliczamy indeksy klatek do zachowania
|
|
indices = np.linspace(0, len(source_frames) - 1, target_count)
|
|
indices = np.round(indices).astype(int)
|
|
|
|
# Wybieramy klatki
|
|
selected_frames = [source_frames[i] for i in indices]
|
|
|
|
print(f"[LiquidTime] Selected frame indices: {indices}")
|
|
return selected_frames
|
|
|
|
def create_time_mapping(self, num_points, start_curve, middle_curve, middle_moment, end_curve):
|
|
"""
|
|
Ulepszone mapowanie czasu z mocniejszym wpływem parametrów.
|
|
"""
|
|
print(f"[LiquidTime] Creating enhanced time mapping with curves:")
|
|
print(f" - Start curve: {start_curve}")
|
|
print(f" - Middle curve: {middle_curve} at {middle_moment}%")
|
|
print(f" - End curve: {end_curve}")
|
|
|
|
# Konwertujemy middle_moment z procentów na zakres 0-1
|
|
middle_t = middle_moment / 100.0
|
|
|
|
# Generujemy bazowe punkty czasowe
|
|
t_values = np.linspace(0, 1, num_points)
|
|
|
|
# Wzmocnienie efektu krzywych
|
|
curve_power = 2.5 # Zwiększa intensywność efektu
|
|
|
|
def enhanced_curve_influence(value):
|
|
"""Wzmacnia wpływ krzywej zachowując znak"""
|
|
return np.sign(value) * abs(value / 100.0) ** 0.75 * curve_power
|
|
|
|
# Konwertujemy wartości krzywych
|
|
start_influence = enhanced_curve_influence(start_curve)
|
|
middle_influence = enhanced_curve_influence(middle_curve)
|
|
end_influence = enhanced_curve_influence(end_curve)
|
|
|
|
def transform_time(t):
|
|
# Bazowa transformacja dla middle_moment (aktywna nawet bez middle_curve)
|
|
middle_shift = (middle_t - 0.5) * 0.5 # Maksymalne przesunięcie o 25% w każdą stronę
|
|
t_shifted = t + middle_shift * np.sin(np.pi * t) # Sinusoidalne przesunięcie
|
|
|
|
# Obliczamy wpływy poszczególnych krzywych
|
|
start_factor = np.exp(-4 * t_shifted) * start_influence
|
|
end_factor = np.exp(-4 * (1 - t_shifted)) * end_influence
|
|
|
|
# Wpływ middle_curve z gaussowskim rozkładem wokół middle_moment
|
|
middle_width = 0.3 # Szerszy zakres wpływu
|
|
middle_gaussian = np.exp(-((t_shifted - middle_t) ** 2) / (2 * middle_width ** 2))
|
|
middle_factor = middle_gaussian * middle_influence
|
|
|
|
# Łączymy wszystkie wpływy
|
|
combined_shift = (start_factor + middle_factor + end_factor)
|
|
|
|
# Aplikujemy przesunięcie z zachowaniem granic
|
|
transformed_t = t_shifted + combined_shift
|
|
|
|
# Zapewniamy monotoniczność (zawsze do przodu)
|
|
transformed_t = t + (transformed_t - t) * 0.7 # 0.7 to siła efektu
|
|
|
|
return transformed_t
|
|
|
|
# Aplikujemy transformację do wszystkich punktów
|
|
transformed_points = np.array([transform_time(t) for t in t_values])
|
|
|
|
# Normalizujemy do zakresu [0, 1] zachowując proporcje
|
|
transformed_points = (transformed_points - transformed_points.min()) / \
|
|
(transformed_points.max() - transformed_points.min())
|
|
|
|
# Wygładzamy przejścia z zachowaniem kształtu
|
|
def smooth_points(points, window=5):
|
|
smoothed = np.copy(points)
|
|
half_window = window // 2
|
|
|
|
for i in range(half_window, len(points) - half_window):
|
|
# Ważone średnia z większym wpływem punktów centralnych
|
|
weights = np.exp(-0.5 * (np.arange(window) - half_window) ** 2)
|
|
window_slice = points[i - half_window:i + half_window + 1]
|
|
smoothed[i] = np.sum(window_slice * weights) / np.sum(weights)
|
|
|
|
# Zachowujemy początek i koniec
|
|
smoothed[:half_window] = points[:half_window]
|
|
smoothed[-half_window:] = points[-half_window:]
|
|
return smoothed
|
|
|
|
smoothed_points = smooth_points(transformed_points)
|
|
|
|
# Końcowa normalizacja
|
|
smoothed_points = (smoothed_points - smoothed_points[0]) / \
|
|
(smoothed_points[-1] - smoothed_points[0])
|
|
|
|
# Debug info
|
|
print(f"[LiquidTime] Time mapping statistics:")
|
|
print(f" - Average speed: {1/(len(smoothed_points)-1):.4f}")
|
|
print(f" - Max speed: {np.max(np.diff(smoothed_points)):.4f}")
|
|
print(f" - Min speed: {np.min(np.diff(smoothed_points)):.4f}")
|
|
print(f" - Middle point actual position: {np.argmin(np.abs(smoothed_points - 0.5))/len(smoothed_points)*100:.1f}%")
|
|
|
|
return smoothed_points
|
|
|
|
def interpolate_pair(self, img1, img2, time_step):
|
|
"""Interpolate between two frames"""
|
|
# Convert images to torch tensors in the correct format
|
|
img1_tensor = torch.from_numpy(img1).permute(2, 0, 1).unsqueeze(0)
|
|
img2_tensor = torch.from_numpy(img2).permute(2, 0, 1).unsqueeze(0)
|
|
|
|
# Move to device and cast to correct precision
|
|
img1_tensor = img1_tensor.to(device=self.device, dtype=self.precision)
|
|
img2_tensor = img2_tensor.to(device=self.device, dtype=self.precision)
|
|
|
|
dt = img1_tensor.new_full((1, 1), float(time_step))
|
|
|
|
with torch.no_grad():
|
|
try:
|
|
interpolated = self.model(img1_tensor, img2_tensor, dt)
|
|
return interpolated.squeeze(0).permute(1, 2, 0).cpu().numpy()
|
|
except RuntimeError as e:
|
|
if "out of memory" in str(e):
|
|
torch.cuda.empty_cache()
|
|
print("[LiquidTime] GPU out of memory error occurred!")
|
|
raise RuntimeError("GPU out of memory. Try reducing image size or using CPU mode.")
|
|
raise
|
|
|
|
def interpolate_frames(self, images, target_frames, start_curve=0, middle_curve=0, middle_moment=50, end_curve=0):
|
|
print(f"\n[LiquidTime] Starting interpolation process...")
|
|
print(f"[LiquidTime] Input: {len(images)} frames, Target: {target_frames} frames")
|
|
print(f"[LiquidTime] Curve parameters: start={start_curve}, middle={middle_curve} at {middle_moment}%, end={end_curve}")
|
|
|
|
# Konwertujemy input na numpy jeśli potrzeba
|
|
if torch.is_tensor(images):
|
|
images = images.cpu().numpy()
|
|
|
|
num_source_frames = len(images)
|
|
|
|
# Obsługa redukcji klatek
|
|
if target_frames < num_source_frames:
|
|
print(f"[LiquidTime] Reducing frames from {num_source_frames} to {target_frames}")
|
|
images = self.sample_frames(images, target_frames)
|
|
return (torch.from_numpy(np.stack(images)),)
|
|
|
|
# Obsługa równej liczby klatek
|
|
if target_frames == num_source_frames:
|
|
print(f"[LiquidTime] Target frames equals source frames, returning without changes")
|
|
return (torch.from_numpy(np.stack(images)),)
|
|
|
|
# Interpolacja
|
|
self.load_model()
|
|
|
|
# Obliczamy punkty czasowe używając krzywej czasowej
|
|
time_points = self.create_time_mapping(target_frames, start_curve, middle_curve, middle_moment, end_curve)
|
|
|
|
# Obliczamy ile klatek przypada na każdy segment źródłowy
|
|
source_segments = np.linspace(0, 1, num_source_frames)
|
|
output_frames = []
|
|
|
|
try:
|
|
# Main interpolation loop with progress bar
|
|
segment_progress = 0
|
|
last_progress_update = 0
|
|
progress_update_interval = max(1, target_frames // 100) # Aktualizacja co 1% postępu
|
|
|
|
with tqdm(total=target_frames,
|
|
desc="[LiquidTime] Generating frames",
|
|
bar_format='{desc}: {percentage:3.0f}%|{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}]') as pbar:
|
|
|
|
current_source_idx = 0
|
|
for frame_idx, t in enumerate(time_points):
|
|
# Znajdujemy odpowiedni segment źródłowy
|
|
while (current_source_idx < len(source_segments) - 1 and
|
|
t > source_segments[current_source_idx + 1]):
|
|
current_source_idx += 1
|
|
segment_progress = 0
|
|
|
|
if current_source_idx >= len(images) - 1:
|
|
current_source_idx = len(images) - 2
|
|
|
|
# Obliczamy lokalny time_step w aktualnym segmencie
|
|
local_segment_start = source_segments[current_source_idx]
|
|
local_segment_end = source_segments[current_source_idx + 1]
|
|
local_t = (t - local_segment_start) / (local_segment_end - local_segment_start)
|
|
local_t = max(0, min(1, local_t)) # Clamp to [0, 1]
|
|
|
|
# Obliczamy procent ukończenia aktualnego segmentu
|
|
segment_progress = local_t * 100
|
|
|
|
# Aktualizujemy informacje o postępie
|
|
if frame_idx - last_progress_update >= progress_update_interval:
|
|
overall_progress = (frame_idx / target_frames) * 100
|
|
print(f"[LiquidTime] Processing: Segment {current_source_idx + 1}/{len(images) - 1} "
|
|
f"({segment_progress:.1f}% complete) | "
|
|
f"Overall: {overall_progress:.1f}% | "
|
|
f"Frames: {frame_idx}/{target_frames}")
|
|
last_progress_update = frame_idx
|
|
|
|
# Interpolujemy klatkę
|
|
interpolated = self.interpolate_pair(
|
|
images[current_source_idx],
|
|
images[current_source_idx + 1],
|
|
local_t
|
|
)
|
|
output_frames.append(interpolated)
|
|
pbar.update(1)
|
|
|
|
# Informacja o rozpoczęciu nowego segmentu
|
|
if segment_progress == 0:
|
|
print(f"[LiquidTime] Starting new segment: {current_source_idx + 1}/{len(images) - 1}")
|
|
|
|
# Okresowe czyszczenie pamięci GPU
|
|
if frame_idx % 100 == 0:
|
|
torch.cuda.empty_cache()
|
|
|
|
final_frame_count = len(output_frames)
|
|
print(f"[LiquidTime] Generated frames breakdown:")
|
|
print(f" - Initial source frames: {num_source_frames}")
|
|
print(f" - Total output frames: {final_frame_count}")
|
|
print(f" - Target frames: {target_frames}")
|
|
|
|
if final_frame_count != target_frames:
|
|
print(f"[LiquidTime] WARNING: Frame count mismatch!")
|
|
print(f" Expected: {target_frames}, Got: {final_frame_count}")
|
|
else:
|
|
print(f"[LiquidTime] Frame count matches target exactly!")
|
|
|
|
# Convert back to torch tensor for ComfyUI compatibility
|
|
result = torch.from_numpy(np.stack(output_frames))
|
|
print(f"[LiquidTime] Output shape: {result.shape}")
|
|
print(f"[LiquidTime] Process completed successfully!")
|
|
|
|
# Dodajemy informacje o autorze i wsparciu
|
|
print("\n" + "="*50)
|
|
print("LiquidTime - by PabloGFX")
|
|
print("YouTube: https://youtube.com/@lazniak")
|
|
print("For support: https://buymeacoffee.com/eyb8tkx3to")
|
|
print("="*50 + "\n")
|
|
|
|
return (result,)
|
|
|
|
except Exception as e:
|
|
print(f"[LiquidTime] Error during interpolation: {str(e)}")
|
|
raise
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"LiquidTime": LiquidTimeNode
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"LiquidTime": "LiquidTime - by PabloGFX"
|
|
}
|