# Funciones que ayudarñán en el resto de la app

import numpy as mp
import torch
from huggan.pytorch.lightweigth_gan.lightweigth_gan import  lightweightGAN #par generar imágenes

## Cargamos el modelo desde el Hub de Hugging Face
def carga_modelo(model_name="ceyda/butterfly_cropped_uniq1K_512", model_version=None):
    #creamos nuestra gan y descargamos el modelo
    gan = LightweightGAN.from_pretrained(model_name, version=model_version)
    #ponemos nuestra gan en forma de evaluacion
    gan.eval()
    return gan


## Usamos el modelo GAN para generar imágenes, y le decimos que el tamaño será 1
def genera(gan, batch_size=1):
    #utilizando pytorch vamos a llamar a nuestra gan 
    with torch.no_grad():
        ims = gan.G(torch.randn(batch_size, gan.latent_dim)).clamp_(0.0, 1.0) * 255
        ims = ims.permute(0, 2, 3, 1).detach().cpu().numpy().astype(np.uint8)
    return ims