Update app.py

app.py

@@ -2,29 +2,22 @@ import gradio as gr
 from transformers import AutoProcessor, AutoModelForCausalLM
 import spaces
 import torch.nn.functional as F
-import requests
 import copy
 import torch
-from PIL import Image, ImageDraw, ImageFont
-import io
-import matplotlib.pyplot as plt
-import matplotlib.patches as patches
 
 import random
 import numpy as np
 from esm import pretrained, FastaBatchedDataset
 
 
+def get_model(model_id):
+    a, b = pretrained.load_model_and_alphabet(model_id.split('/')[1])
+    return (a, b)
+
 models = {
-    'facebook/esm2_t36_3B_UR50D': pretrained.load_model_and_alphabet('esm2_t36_3B_UR50D').to("cuda").eval(),
+    'facebook/esm2_t36_3B_UR50D': get_model('facebook/esm2_t36_3B_UR50D'),
     }
 
-processors = {
-    'microsoft/Florence-2-large-ft': AutoProcessor.from_pretrained('microsoft/Florence-2-large-ft', trust_remote_code=True),
-    'microsoft/Florence-2-large': AutoProcessor.from_pretrained('microsoft/Florence-2-large', trust_remote_code=True),
-    'microsoft/Florence-2-base-ft': AutoProcessor.from_pretrained('microsoft/Florence-2-base-ft', trust_remote_code=True),
-    'microsoft/Florence-2-base': AutoProcessor.from_pretrained('microsoft/Florence-2-base', trust_remote_code=True),
-}
 
 
 DESCRIPTION = "Esm2 embedding"