Spaces:
Build error
Build error
| import torch | |
| import torch.nn.functional as F | |
| import gradio as gr | |
| from PIL import Image | |
| import pickle | |
| import json | |
| from huggingface_hub import hf_hub_download | |
| from model import image_captioning_model | |
| from tokenizer import CaptionTokenizer # Важно! | |
| # Загрузка файлов | |
| repo_id = "Pumponchik/ImageCaption" | |
| repo_type = "space" | |
| model_path = hf_hub_download(repo_id=repo_id, filename="pytorch_model.bin", repo_type=repo_type) | |
| config_path = hf_hub_download(repo_id=repo_id, filename="config.json", repo_type=repo_type) | |
| tokenizer_path = hf_hub_download(repo_id=repo_id, filename="tokenizer.pkl", repo_type=repo_type) | |
| transform_path = hf_hub_download(repo_id=repo_id, filename="image_transforms.pkl", repo_type=repo_type) | |
| # Загрузка конфигурации | |
| with open(config_path, "r") as f: | |
| full_config = json.load(f) | |
| # Фильтрация параметров для модели | |
| model_params = { | |
| "vocab_size": full_config["vocab_size"], | |
| "embedding_dim": full_config["embedding_dim"], | |
| "hidden_size": full_config["hidden_size"], | |
| "num_layers": full_config["num_layers"] | |
| } | |
| # Инициализация модели | |
| model = image_captioning_model(**model_params) | |
| # Загрузка весов модели | |
| model.load_state_dict(torch.load(model_path, map_location="cpu"), strict=False) | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| model.to(device) | |
| model.eval() | |
| tokenizer = CaptionTokenizer.load(tokenizer_path) | |
| # Загрузка трансформаций | |
| with open(transform_path, "rb") as f: | |
| transform = pickle.load(f) | |
| def wrap_text(text, max_width): | |
| words = text.split(" ") | |
| result = [[words[0]]] | |
| for word in words[1:]: | |
| if len(" ".join(result[-1])) + len(word) + 1 > max_width: | |
| result[-1] = " ".join(result[-1]) | |
| result.append([]) | |
| result[-1].append(word) | |
| result[-1] = " ".join(result[-1]) | |
| return "\n".join(result) | |
| def generate(model, image, tokenizer, max_seq_len=20, top_p=None, top_k=1): | |
| assert top_p is None or top_k is None, "Не используйте одновременно top_p и top_k" | |
| model.eval() | |
| device = next(model.parameters()).device | |
| image = transform(image).to(device).unsqueeze(0) | |
| batch_size = image.size(0) | |
| with torch.no_grad(): | |
| img_emb = model.img_fe(image) | |
| h_0 = model.img_to_hidden(img_emb).unsqueeze(0).repeat(model.text_fe.num_layers, 1, 1) | |
| c_0 = h_0.clone() | |
| input_token = torch.full((batch_size,), tokenizer.tok_to_ind['<BOS>'], dtype=torch.long, device=device) | |
| result_tokens = [[] for _ in range(batch_size)] | |
| mask = torch.ones(batch_size, dtype=torch.bool, device=device) | |
| for _ in range(max_seq_len): | |
| input_emb = model.text_fe.embedding(input_token).unsqueeze(1) | |
| output, (h_0, c_0) = model.text_fe.rnn(input_emb, (h_0, c_0)) | |
| logits = model.fc_out(output.squeeze(1)) | |
| probs = F.softmax(logits, dim=-1) | |
| next_tokens = torch.zeros(batch_size, dtype=torch.long, device=device) | |
| for i in range(batch_size): | |
| if not mask[i]: | |
| continue | |
| prob = probs[i] | |
| if top_k is not None: | |
| sorted_probs, sorted_indices = torch.sort(prob, descending=True) | |
| sorted_probs = sorted_probs[:top_k] | |
| sorted_indices = sorted_indices[:top_k] | |
| sorted_probs = sorted_probs / sorted_probs.sum() | |
| next_token_idx = torch.multinomial(sorted_probs, 1) | |
| next_tokens[i] = sorted_indices[next_token_idx] | |
| elif top_p is not None: | |
| sorted_probs, sorted_indices = torch.sort(prob, descending=True) | |
| cumulative_probs = torch.cumsum(sorted_probs, dim=0) | |
| mask_p = cumulative_probs <= top_p | |
| sorted_probs = sorted_probs[mask_p] | |
| sorted_indices = sorted_indices[mask_p] | |
| sorted_probs = sorted_probs / sorted_probs.sum() | |
| next_token_idx = torch.multinomial(sorted_probs, 1) | |
| next_tokens[i] = sorted_indices[next_token_idx] | |
| else: | |
| next_tokens[i] = torch.multinomial(prob, 1) | |
| token_id = next_tokens[i].item() | |
| result_tokens[i].append(token_id) | |
| if token_id == tokenizer.tok_to_ind['<EOS>']: | |
| mask[i] = False | |
| input_token = next_tokens | |
| if not mask.any(): | |
| break | |
| result_text = [] | |
| for tokens in result_tokens: | |
| caption = " ".join([tokenizer.ind_to_tok.get(t, "<UNK>") for t in tokens | |
| if t not in [tokenizer.tok_to_ind.get("<BOS>", -1), | |
| tokenizer.tok_to_ind.get("<EOS>", -1), | |
| tokenizer.tok_to_ind.get("<PAD>", -1)]]) | |
| result_text.append(caption) | |
| return result_text[0] if batch_size == 1 else result_text | |
| def generate_caption(image): | |
| image = image.convert("RGB") | |
| max_seq_len = full_config.get("max_len", 20) | |
| caption = generate(model, image, tokenizer, max_seq_len=max_seq_len, top_k=5) | |
| caption = wrap_text(caption, 43) | |
| return caption | |
| iface = gr.Interface( | |
| fn=generate_caption, | |
| inputs=gr.Image(type="pil", label="Загрузите изображение"), | |
| outputs=gr.Textbox(label="Сгенерированное описание"), | |
| title="Генерация описаний изображений", | |
| description="Загрузите изображение, и модель создаст текстовое описание." | |
| ) | |
| iface.launch( | |
| server_name="0.0.0.0", | |
| server_port=7860, | |
| ) | |