Spaces:
Runtime error
Runtime error
import torch | |
from huggan.pytorch.lightweight_gan.lightweight_gan import LightweightGAN | |
from datasets import load_dataset | |
def get_train_data(dataset_name="ceyda/smithsonian_butterflies_transparent_cropped",data_limit=1000): | |
dataset=load_dataset(dataset_name) | |
dataset=dataset.sort("sim_score") | |
score_thresh = dataset["train"][data_limit]['sim_score'] | |
dataset = dataset.filter(lambda x: x['sim_score'] < score_thresh) | |
dataset = dataset.map(lambda x: {'image' : x['image'].convert("RGB")}) | |
return dataset["train"] | |
from transformers import BeitFeatureExtractor, BeitForImageClassification | |
feature_extractor = BeitFeatureExtractor.from_pretrained('microsoft/beit-base-patch16-224') | |
model = BeitForImageClassification.from_pretrained('microsoft/beit-base-patch16-224') | |
def embed(images): | |
inputs = feature_extractor(images=images, return_tensors="pt") | |
outputs = model(**inputs,output_hidden_states= True) | |
last_hidden=outputs.hidden_states[-1] | |
pooler=model.base_model.pooler | |
final_emb=pooler(last_hidden).detach().numpy() | |
return final_emb | |
def build_index(): | |
dataset=get_train_data() | |
ds_with_embeddings = dataset.map(lambda x: {"beit_embeddings":embed(x["image"])},batched=True,batch_size=20) | |
ds_with_embeddings.add_faiss_index(column='beit_embeddings') | |
ds_with_embeddings.save_faiss_index('beit_embeddings', 'beit_index.faiss') | |
def get_dataset(): | |
dataset=get_train_data() | |
dataset.load_faiss_index('beit_embeddings', 'beit_index.faiss') | |
return dataset | |
def load_model(model_name='ceyda/butterfly_cropped_uniq1K_512'): | |
gan = LightweightGAN.from_pretrained(model_name) | |
gan.eval() | |
return gan | |
def generate(gan,batch_size=1): | |
with torch.no_grad(): | |
ims = gan.G(torch.randn(batch_size, gan.latent_dim)).clamp_(0., 1.) | |
ims = ims.permute(0,2,3,1).detach().cpu().numpy() | |
return ims | |
def interpolate(): | |
pass |