Update utils.py

utils.py

@@ -6,72 +6,109 @@ import torch
 from PIL import Image
 
 def parse_weight_string(string, max_frames):
+    """
+    Parses complex Deforum weight strings with math support (sin, cos, t).
+    """
     string = re.sub(r'\s+', '', str(string))
     keyframes = {}
     parts = string.split(',')
+    
     for part in parts:
         try:
             if ':' not in part: continue
             f_str, v_str = part.split(':', 1)
             keyframes[int(f_str)] = v_str.strip('()')
         except: continue
+        
     if 0 not in keyframes: keyframes[0] = "0"
     
     series = np.zeros(int(max_frames))
-    sorted_k = sorted(keyframes.keys())
-    for i in range(len(sorted_k)):
-        f_start = sorted_k[i]
-        f_end = sorted_k[i+1] if i < len(sorted_k)-1 else int(max_frames)
+    sorted_keys = sorted(keyframes.keys())
+    
+    for i in range(len(sorted_keys)):
+        f_start = sorted_keys[i]
+        f_end = sorted_keys[i+1] if i < len(sorted_keys)-1 else int(max_frames)
         formula = keyframes[f_start]
+        
         for f in range(f_start, f_end):
+            t = f
             try:
-                series[f] = float(numexpr.evaluate(formula, local_dict={'t':f,'pi':np.pi,'sin':np.sin,'cos':np.cos}))
+                val = numexpr.evaluate(formula, local_dict={'t': t, 'pi': np.pi, 'sin': np.sin, 'cos': np.cos, 'tan': np.tan, 'abs': np.abs})
+                series[f] = float(val)
             except:
-                series[f] = float(formula) if formula.replace('.','',1).isdigit() else (series[f-1] if f>0 else 0.0)
+                try: series[f] = float(formula)
+                except: series[f] = series[f-1] if f > 0 else 0.0
+                
     return series
 
-def interpolate_prompts(pipe, prompt_dict, max_frames):
-    """Blends CLIP embeddings between keyframes for smooth conceptual transitions."""
-    sorted_keys = sorted(prompt_dict.keys())
-    # Pre-calculate embeddings for all keyframe prompts
-    key_embs = {}
-    for k in sorted_keys:
-        tokens = pipe.tokenizer(prompt_dict[k], padding="max_length", max_length=pipe.tokenizer.model_max_length, truncation=True, return_tensors="pt").input_ids.to(pipe.device)
-        with torch.no_grad():
-            key_embs[k] = pipe.text_encoder(tokens)[0]
-    
-    full_embs = []
-    for f in range(max_frames):
-        # Find surrounding keyframes
-        before = [k for k in sorted_keys if k <= f]
-        after = [k for k in sorted_keys if k > f]
-        
-        if not after:
-            full_embs.append(key_embs[before[-1]])
-        elif not before:
-            full_embs.append(key_embs[after[0]])
-        else:
-            k1, k2 = before[-1], after[0]
-            alpha = (f - k1) / (k2 - k1)
-            # Spherical Linear Interpolation (Slerp) or simple Lerp
-            blended = torch.lerp(key_embs[k1], key_embs[k2], alpha)
-            full_embs.append(blended)
-    return full_embs
+def get_border_mode(mode_str):
+    return {
+        'Reflect': cv2.BORDER_REFLECT_101,
+        'Replicate': cv2.BORDER_REPLICATE,
+        'Wrap': cv2.BORDER_WRAP,
+        'Black': cv2.BORDER_CONSTANT
+    }.get(mode_str, cv2.BORDER_REFLECT_101)
 
-def maintain_colors(img, anchor, mode='LAB'):
-    if mode == 'None' or anchor is None: return img
-    img_np, anc_np = np.array(img), np.array(anchor)
+def maintain_colors(image, anchor, mode='LAB'):
+    """
+    Matches the color distribution of 'image' to 'anchor'.
+    """
+    if mode == 'None' or anchor is None: return image
+    
+    img_np = np.array(image).astype(np.uint8)
+    anc_np = np.array(anchor).astype(np.uint8)
+    
     if mode == 'LAB':
-        img_lab, anc_lab = cv2.cvtColor(img_np, cv2.COLOR_RGB2LAB), cv2.cvtColor(anc_np, cv2.COLOR_RGB2LAB)
+        img_cvt = cv2.cvtColor(img_np, cv2.COLOR_RGB2LAB)
+        anc_cvt = cv2.cvtColor(anc_np, cv2.COLOR_RGB2LAB)
         for i in range(3):
-            img_lab[:,:,i] = np.clip(img_lab[:,:,i] - np.mean(img_lab[:,:,i]) + np.mean(anc_lab[:,:,i]), 0, 255)
-        return Image.fromarray(cv2.cvtColor(img_lab, cv2.COLOR_LAB2RGB))
-    return img
+            img_cvt[:,:,i] = np.clip(img_cvt[:,:,i] - img_cvt[:,:,i].mean() + anc_cvt[:,:,i].mean(), 0, 255)
+        out = cv2.cvtColor(img_cvt, cv2.COLOR_LAB2RGB)
+        
+    elif mode == 'HSV':
+        img_cvt = cv2.cvtColor(img_np, cv2.COLOR_RGB2HSV)
+        anc_cvt = cv2.cvtColor(anc_np, cv2.COLOR_RGB2HSV)
+        # Match S and V, keep Hue
+        for i in [1, 2]:
+            img_cvt[:,:,i] = np.clip(img_cvt[:,:,i] - img_cvt[:,:,i].mean() + anc_cvt[:,:,i].mean(), 0, 255)
+        out = cv2.cvtColor(img_cvt, cv2.COLOR_HSV2RGB)
+        
+    elif mode == 'RGB':
+        for i in range(3):
+            img_np[:,:,i] = np.clip(img_np[:,:,i] - img_np[:,:,i].mean() + anc_np[:,:,i].mean(), 0, 255)
+        out = img_np
+        
+    else:
+        return image
+
+    return Image.fromarray(out)
 
-def anim_frame_warp_2d(img, args, mode='Reflect'):
-    cv_img = np.array(img)
+def anim_frame_warp_2d(prev_img, args, border_mode_str):
+    """
+    Applies 2D affine transformation (Zoom, Rotate, Pan).
+    """
+    if prev_img is None: return None
+    cv_img = np.array(prev_img)
     h, w = cv_img.shape[:2]
-    mat = cv2.getRotationMatrix2D((w//2, h//2), args.get('angle',0), args.get('zoom',1))
-    mat[0, 2] += args.get('tx',0); mat[1, 2] += args.get('ty',0)
-    b = {'Reflect':cv2.BORDER_REFLECT_101, 'Replicate':cv2.BORDER_REPLICATE, 'Wrap':cv2.BORDER_WRAP}[mode]
-    return Image.fromarray(cv2.warpAffine(cv_img, mat, (w, h), borderMode=b))
+    center = (w // 2, h // 2)
+    
+    angle = args.get('angle', 0)
+    zoom = args.get('zoom', 1.0)
+    tx = args.get('tx', 0)
+    ty = args.get('ty', 0)
+    
+    # Create Matrix
+    mat = cv2.getRotationMatrix2D(center, angle, zoom)
+    mat[0, 2] += tx
+    mat[1, 2] += ty
+    
+    border = get_border_mode(border_mode_str)
+    warped = cv2.warpAffine(cv_img, mat, (w, h), borderMode=border)
+    return Image.fromarray(warped)
+
+def add_noise(img, noise_amt):
+    if noise_amt <= 0: return img
+    img_np = np.array(img).astype(np.float32)
+    noise = np.random.normal(0, noise_amt * 255, img_np.shape).astype(np.float32)
+    noisy = np.clip(img_np + noise, 0, 255).astype(np.uint8)
+    return Image.fromarray(noisy)