Update app.py

app.py

@@ -56,8 +56,6 @@ model2.config.id2label = mapping
 model3 = models.resnet101(pretrained=False)
 model3.fc = torch.nn.Linear(2048, len(genres))
 
-device = torch.device('cpu')
-
 
 # In[5]:
 
@@ -94,19 +92,19 @@ class Multimodal(torch.nn.Module):
 
 model = Multimodal(model1, model2, model3)
 model.load_state_dict(torch.load('multimodel.pt', map_location=torch.device('cpu')))
-model.to(device)
 model.eval()
-model_gen.to(device)
-model_gen.eval()
+device = torch.device('cpu')
 
 
 # In[8]:
 
 
-def generate_plot(title: str, model: AutoModelForSeq2SeqLM, tokenizer: AutoTokenizer) -> str:
+def generate_plot(title: str, model: AutoModelForSeq2SeqLM, tokenizer: AutoTokenizer, device) -> str:
     quote = 'What is the story of the movie {}?'
+    model_gen.to(device)
+    model_gen.eval()
 
-    input_ids = tokenizer(quote.format(title), return_tensors='pt')
+    input_ids = tokenizer(quote.format(title), return_tensors='pt').input_ids.to(device)
     output = model.generate(input_ids, max_length=256, do_sample=True, temperature=0.09)
     return tokenizer.decode(output[0], skip_special_tokens=True)
 
@@ -117,24 +115,25 @@ def generate_plot(title: str, model: AutoModelForSeq2SeqLM, tokenizer: AutoToken
 def inference(title, image, 
               tokenizer1=tokenizer1, tokenizer2=tokenizer2, tokenizer_gen=tokenizer_gen,
               model_gen=model_gen, model=model, 
-              genres=genres):
+              genres=genres, device=device):
     title_input = tokenizer1(title, return_tensors='pt', padding=True, truncation=True)
-    title_input_ids = title_input['input_ids']
-    title_attention_mask = title_input['attention_mask']
+    title_input_ids = title_input['input_ids'].to(device)
+    title_attention_mask = title_input['attention_mask'].to(device)
 
-    plot = generate_plot(title, model_gen, tokenizer_gen)
+    plot = generate_plot(title, model_gen, tokenizer_gen, device)
     plot_input = tokenizer2(plot, return_tensors='pt', padding=True, truncation=True)
-    plot_input_ids = plot_input['input_ids']
-    plot_attention_mask = plot_input['attention_mask']
+    plot_input_ids = plot_input['input_ids'].to(device)
+    plot_attention_mask = plot_input['attention_mask'].to(device)
 
     # If image is not uploaded
     if image is None:
-        image_input = torch.zeros((1, 3, 224, 224))
+        image_input = torch.zeros((1, 3, 224, 224)).to(device)
 
     else:
         image_input = image.resize((224, 224))
         image_input = v2.ToTensor()(image_input)
         image_input = image_input.unsqueeze(0)
+        image_input = image_input.to(device)
 
     output = model(title_input_ids, title_attention_mask, plot_input_ids, plot_attention_mask, image_input)
     output = torch.sigmoid(output)
@@ -163,3 +162,4 @@ app = gr.Interface(fn=inference, inputs=["text", "pil"], outputs="text", title="
 
 
 app.launch(share=True)
+