diff --git a/src/inferencescale/utils.py b/src/inferencescale/utils.py new file mode 100644 index 0000000..a5907de --- /dev/null +++ b/src/inferencescale/utils.py @@ -0,0 +1,44 @@ +import gc +import torch +import torch.nn.functional as F + +import time +import numpy as np +import comfy.model_management as model_management + + +def clear_memory(model_name): + """ + Unloads the specified model from memory (if loaded) and clears PyTorch cache. + """ + print("Unloading Model:") + loaded_models = model_management.loaded_models() + + if model_name in loaded_models: + print(" - Model found in memory, unloading...") + loaded_models.remove(model_name) + + model_management.free_memory(1e30, model_management.get_torch_device(), loaded_models) + model_management.soft_empty_cache(True) + + # Clear Python and CUDA caches + try: + print("Clearing Cache...") + gc.collect() + torch.cuda.empty_cache() + torch.cuda.ipc_collect() + except: + print("Unable to clear cache") + time.sleep(2) + + +def generate_neighbors(x, threshold=0.95, num_neighbors=4): + """Generates neighbors on the same sphere with cosine similarity = threshold.""" + # Function adopted from: https://github.com/sayakpaul/tt-scale-flux + rng = np.random.Generator(np.random.PCG64()) + x_f = x.flatten(1) + x_norm = torch.linalg.norm(x_f, dim=-1, keepdim=True, dtype=torch.float64).unsqueeze(-2) + u = x_f.unsqueeze(-2) / x_norm.clamp_min(1e-12) + v = torch.from_numpy(rng.standard_normal(size=(u.shape[0], num_neighbors, u.shape[-1]), dtype=np.float64)).to(u.device) + w = F.normalize(v - (v @ u.transpose(-2, -1)) * u, dim=-1) + return (x_norm * (threshold * u + np.sqrt(1 - threshold**2) * w)).reshape(x.shape[0], num_neighbors, *x.shape[1:]).to(x.dtype)