# Copyright (c) 2015-present, Facebook, Inc.
# All rights reserved.
#
# This source code is licensed under the CC-by-NC license found in the
# LICENSE file in the root directory of this source tree.
#
import torch


def coeffs2text(coeffs, tokenizer):
    learned_ids = coeffs.argmax(dim=-1).tolist()[0]
    text = tokenizer.decode(learned_ids)
    return text


def get_scale(pipe, device, t=None):
    from diffusers import DDPMScheduler

    model_id = f"CompVis/stable-diffusion-v1-4"
    cache_dir = '../../cache'
    # device = f'cuda:0'
    torch_dtype = torch.float16
    # tokenizer, text_encoder, unet, vae, pipe = get_stable_diffusion_pipeline(model_id, torch_dtype, cache_dir, device, True)

    num_inference_steps = 50
    pipe.scheduler.set_timesteps(num_inference_steps, device=device)
    timesteps = pipe.scheduler.timesteps


    noise_scheduler = DDPMScheduler.from_pretrained(model_id, subfolder="scheduler")

    alphas_cumprod = noise_scheduler.alphas_cumprod.to(device=device, dtype=torch_dtype)
    sqrt_alpha_prod = alphas_cumprod[timesteps] ** 0.5
    sqrt_alpha_prod = sqrt_alpha_prod.flatten()

    sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[timesteps]) ** 0.5
    sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten()


    alphas = noise_scheduler.alphas.to(device=device, dtype=torch_dtype)
    betas = noise_scheduler.betas.to(device=device, dtype=torch_dtype)
    var = betas
    scales = betas**2 / (2 * var * alphas * (1 - alphas_cumprod))

    if t is None:
        return scales
    else:
        return scales[t]
