Spaces:
Running
Running
| from time import gmtime, strftime | |
| print(f'{strftime("%Y-%m-%d %H:%M:%S", gmtime())} Preparing for inference...') # noqa | |
| from rudalle.pipelines import generate_images | |
| from rudalle import get_rudalle_model, get_tokenizer, get_vae | |
| from huggingface_hub import hf_hub_url, cached_download | |
| import torch | |
| from io import BytesIO | |
| import base64 | |
| print(f"GPUs available: {torch.cuda.device_count()}") | |
| print(f"GPU[0] memory: {int(torch.cuda.get_device_properties(0).total_memory / 1048576)}Mib") | |
| print(f"GPU[0] memory reserved: {int(torch.cuda.memory_reserved(0) / 1048576)}Mib") | |
| print(f"GPU[0] memory allocated: {int(torch.cuda.memory_allocated(0) / 1048576)}Mib") | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| fp16 = torch.cuda.is_available() | |
| file_dir = "./models" | |
| file_name = "pytorch_model.bin" | |
| config_file_url = hf_hub_url(repo_id="minimaxir/ai-generated-pokemon-rudalle", filename=file_name) | |
| cached_download(config_file_url, cache_dir=file_dir, force_filename=file_name) | |
| model = get_rudalle_model('Malevich', pretrained=False, fp16=fp16, device=device) | |
| model.load_state_dict(torch.load(f"{file_dir}/{file_name}", map_location=f"{'cuda:0' if torch.cuda.is_available() else 'cpu'}")) | |
| vae = get_vae().to(device) | |
| tokenizer = get_tokenizer() | |
| print(f'{strftime("%Y-%m-%d %H:%M:%S", gmtime())} Ready for inference') | |
| def english_to_russian(english_string): | |
| word_map = { | |
| "grass": "Покемон трава", | |
| "fire": "Покемон огня", | |
| "water": "Покемон в воду", | |
| "lightning": "Покемон электрического типа", | |
| "fighting": "Покемон боевого типа", | |
| "psychic": "Покемон психического типа", | |
| "colorless": "Покемон нормального типа", | |
| "darkness": "Покемон темного типа", | |
| "metal": "Покемон из стали типа", | |
| "dragon": "Покемон типа дракона", | |
| "fairy": "Покемон фея" | |
| } | |
| return word_map[english_string.lower()] | |
| def generate_image(prompt): | |
| if prompt.lower() in ['grass', 'fire', 'water', 'lightning', 'fighting', 'psychic', 'colorless', 'darkness', | |
| 'metal', 'dragon', 'fairy']: | |
| prompt = english_to_russian(prompt) | |
| result, _ = generate_images(prompt, tokenizer, model, vae, top_k=2048, images_num=1, top_p=0.995) | |
| buffer = BytesIO() | |
| result[0].save(buffer, format="PNG") | |
| base64_bytes = base64.b64encode(buffer.getvalue()) | |
| base64_string = base64_bytes.decode("UTF-8") | |
| return "data:image/png;base64," + base64_string | |