|
|
|
|
|
|
|
|
import torch |
|
|
from typing import Optional, Union, List, Any, Dict |
|
|
import numpy as np |
|
|
from diffusers.utils import load_image |
|
|
import random |
|
|
import time |
|
|
|
|
|
def calculate_shift( |
|
|
image_seq_len, |
|
|
base_seq_len: int = 256, |
|
|
max_seq_len: int = 4096, |
|
|
base_shift: float = 0.5, |
|
|
max_shift: float = 1.16, |
|
|
): |
|
|
m = (max_shift - base_shift) / (max_seq_len - base_seq_len) |
|
|
b = base_shift - m * base_seq_len |
|
|
mu = image_seq_len * m + b |
|
|
return mu |
|
|
|
|
|
def retrieve_timesteps( |
|
|
scheduler, |
|
|
num_inference_steps: Optional[int] = None, |
|
|
device: Optional[Union[str, torch.device]] = None, |
|
|
timesteps: Optional[List[int]] = None, |
|
|
sigmas: Optional[List[float]] = None, |
|
|
**kwargs, |
|
|
): |
|
|
if timesteps is not None and sigmas is not None: |
|
|
raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values") |
|
|
if timesteps is not None: |
|
|
scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs) |
|
|
timesteps = scheduler.timesteps |
|
|
num_inference_steps = len(timesteps) |
|
|
elif sigmas is not None: |
|
|
scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs) |
|
|
timesteps = scheduler.timesteps |
|
|
num_inference_steps = len(timesteps) |
|
|
else: |
|
|
scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) |
|
|
timesteps = scheduler.timesteps |
|
|
return timesteps, num_inference_steps |
|
|
|
|
|
|
|
|
def load_image_from_path(image_path: str): |
|
|
"""Loads an image from a given file path.""" |
|
|
return load_image(image_path) |
|
|
|
|
|
def randomize_seed_if_needed(randomize_seed: bool, seed: int, max_seed: int) -> int: |
|
|
"""Randomizes the seed if requested.""" |
|
|
if randomize_seed: |
|
|
return random.randint(0, max_seed) |
|
|
return seed |
|
|
|
|
|
class calculateDuration: |
|
|
def __init__(self, activity_name=""): |
|
|
self.activity_name = activity_name |
|
|
|
|
|
def __enter__(self): |
|
|
self.start_time = time.time() |
|
|
return self |
|
|
|
|
|
def __exit__(self, exc_type, exc_value, traceback): |
|
|
self.end_time = time.time() |
|
|
self.elapsed_time = self.end_time - self.start_time |
|
|
if self.activity_name: |
|
|
print(f"Elapsed time for {self.activity_name}: {self.elapsed_time:.6f} seconds") |
|
|
else: |
|
|
print(f"Elapsed time: {self.elapsed_time:.6f} seconds") |
