Update app.py

app.py

@@ -3,13 +3,13 @@ import gradio as gr
 import torch
 import numpy as np
 from PIL import Image
-from transformers import AutoModel, AutoProcessor
 from torch import nn
 import torch.nn.functional as F
 from datasets import load_dataset
 from torch.utils.data import Dataset, DataLoader
 import os
 from tqdm import tqdm
+from transformers import Florence2Model, Florence2Processor
 
 class SDDataset(Dataset):
     def __init__(self, dataset, processor, model_to_idx, token_to_idx, max_samples=5000):
@@ -49,7 +49,7 @@ class SDRecommenderModel(nn.Module):
         
     def forward(self, image_features):
         # Get Florence embeddings
-        features = self.florence.get_image_features(image_features)
+        features = self.florence.get_image_features(**image_features)
         
         # Generate model and prompt recommendations
         model_logits = self.model_head(features)
@@ -62,16 +62,16 @@ class SDRecommender:
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
         print(f"Using device: {self.device}")
         
-        # Load Florence model and processor dengan trust_remote_code=True
+        # Load Florence model and processor
         print("Loading Florence model and processor...")
-        self.processor = AutoProcessor.from_pretrained(
+        self.processor = Florence2Processor.from_pretrained(
             "microsoft/Florence-2-large",
             trust_remote_code=True
         )
-        self.florence = AutoModel.from_pretrained(
+        self.florence = Florence2Model.from_pretrained(
             "microsoft/Florence-2-large",
             trust_remote_code=True
-        )
+        ).to(self.device)
         
         # Load dataset
         print("Loading dataset...")
@@ -139,7 +139,7 @@ class SDRecommender:
             
             for batch_idx, (images, model_labels, prompt_labels) in enumerate(progress_bar):
                 # Move everything to device
-                images = images.to(self.device)
+                images = {k: v.to(self.device) for k, v in images.items()}
                 model_labels = model_labels.to(self.device)
                 prompt_labels = prompt_labels.to(self.device)
                 
@@ -170,7 +170,8 @@ class SDRecommender:
             image = Image.open(image)
         
         # Process image
-        inputs = self.processor(images=image, return_tensors="pt").to(self.device)
+        inputs = self.processor(images=image, return_tensors="pt")
+        inputs = {k: v.to(self.device) for k, v in inputs.items()}
         
         # Get model predictions
         self.model.eval()