Update app.py

app.py

@@ -9,7 +9,7 @@ from datasets import load_dataset
 from torch.utils.data import Dataset, DataLoader
 import os
 from tqdm import tqdm
-from transformers import Florence2Model, Florence2Processor
+from transformers import AutoProcessor, AutoModelForCausalLM
 
 class SDDataset(Dataset):
     def __init__(self, dataset, processor, model_to_idx, token_to_idx, max_samples=5000):
@@ -44,12 +44,14 @@ class SDRecommenderModel(nn.Module):
     def __init__(self, florence_model, num_models, vocab_size):
         super().__init__()
         self.florence = florence_model
-        self.model_head = nn.Linear(florence_model.config.hidden_size, num_models)
-        self.prompt_head = nn.Linear(florence_model.config.hidden_size, vocab_size)
+        hidden_size = 1024  # Florence-2-large hidden size
+        self.model_head = nn.Linear(hidden_size, num_models)
+        self.prompt_head = nn.Linear(hidden_size, vocab_size)
         
-    def forward(self, image_features):
+    def forward(self, pixel_values):
         # Get Florence embeddings
-        features = self.florence.get_image_features(**image_features)
+        outputs = self.florence(pixel_values=pixel_values, output_hidden_states=True)
+        features = outputs.hidden_states[-1].mean(dim=1)  # Use mean pooling of last hidden state
         
         # Generate model and prompt recommendations
         model_logits = self.model_head(features)
@@ -58,18 +60,19 @@ class SDRecommenderModel(nn.Module):
         return model_logits, prompt_logits
 
 class SDRecommender:
-    def __init__(self, max_samples=1000):
+    def __init__(self, max_samples=500):
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
         print(f"Using device: {self.device}")
         
         # Load Florence model and processor
         print("Loading Florence model and processor...")
-        self.processor = Florence2Processor.from_pretrained(
+        self.processor = AutoProcessor.from_pretrained(
             "microsoft/Florence-2-large",
             trust_remote_code=True
         )
-        self.florence = Florence2Model.from_pretrained(
+        self.florence = AutoModelForCausalLM.from_pretrained(
             "microsoft/Florence-2-large",
+            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
             trust_remote_code=True
         ).to(self.device)