Spaces:
Sleeping
Sleeping
Update app.py
Browse files
app.py
CHANGED
@@ -9,7 +9,7 @@ from datasets import load_dataset
|
|
9 |
from torch.utils.data import Dataset, DataLoader
|
10 |
import os
|
11 |
from tqdm import tqdm
|
12 |
-
from transformers import
|
13 |
|
14 |
class SDDataset(Dataset):
|
15 |
def __init__(self, dataset, processor, model_to_idx, token_to_idx, max_samples=5000):
|
@@ -44,12 +44,14 @@ class SDRecommenderModel(nn.Module):
|
|
44 |
def __init__(self, florence_model, num_models, vocab_size):
|
45 |
super().__init__()
|
46 |
self.florence = florence_model
|
47 |
-
|
48 |
-
self.
|
|
|
49 |
|
50 |
-
def forward(self,
|
51 |
# Get Florence embeddings
|
52 |
-
|
|
|
53 |
|
54 |
# Generate model and prompt recommendations
|
55 |
model_logits = self.model_head(features)
|
@@ -58,18 +60,19 @@ class SDRecommenderModel(nn.Module):
|
|
58 |
return model_logits, prompt_logits
|
59 |
|
60 |
class SDRecommender:
|
61 |
-
def __init__(self, max_samples=
|
62 |
self.device = "cuda" if torch.cuda.is_available() else "cpu"
|
63 |
print(f"Using device: {self.device}")
|
64 |
|
65 |
# Load Florence model and processor
|
66 |
print("Loading Florence model and processor...")
|
67 |
-
self.processor =
|
68 |
"microsoft/Florence-2-large",
|
69 |
trust_remote_code=True
|
70 |
)
|
71 |
-
self.florence =
|
72 |
"microsoft/Florence-2-large",
|
|
|
73 |
trust_remote_code=True
|
74 |
).to(self.device)
|
75 |
|
|
|
9 |
from torch.utils.data import Dataset, DataLoader
|
10 |
import os
|
11 |
from tqdm import tqdm
|
12 |
+
from transformers import AutoProcessor, AutoModelForCausalLM
|
13 |
|
14 |
class SDDataset(Dataset):
|
15 |
def __init__(self, dataset, processor, model_to_idx, token_to_idx, max_samples=5000):
|
|
|
44 |
def __init__(self, florence_model, num_models, vocab_size):
|
45 |
super().__init__()
|
46 |
self.florence = florence_model
|
47 |
+
hidden_size = 1024 # Florence-2-large hidden size
|
48 |
+
self.model_head = nn.Linear(hidden_size, num_models)
|
49 |
+
self.prompt_head = nn.Linear(hidden_size, vocab_size)
|
50 |
|
51 |
+
def forward(self, pixel_values):
|
52 |
# Get Florence embeddings
|
53 |
+
outputs = self.florence(pixel_values=pixel_values, output_hidden_states=True)
|
54 |
+
features = outputs.hidden_states[-1].mean(dim=1) # Use mean pooling of last hidden state
|
55 |
|
56 |
# Generate model and prompt recommendations
|
57 |
model_logits = self.model_head(features)
|
|
|
60 |
return model_logits, prompt_logits
|
61 |
|
62 |
class SDRecommender:
|
63 |
+
def __init__(self, max_samples=500):
|
64 |
self.device = "cuda" if torch.cuda.is_available() else "cpu"
|
65 |
print(f"Using device: {self.device}")
|
66 |
|
67 |
# Load Florence model and processor
|
68 |
print("Loading Florence model and processor...")
|
69 |
+
self.processor = AutoProcessor.from_pretrained(
|
70 |
"microsoft/Florence-2-large",
|
71 |
trust_remote_code=True
|
72 |
)
|
73 |
+
self.florence = AutoModelForCausalLM.from_pretrained(
|
74 |
"microsoft/Florence-2-large",
|
75 |
+
torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
|
76 |
trust_remote_code=True
|
77 |
).to(self.device)
|
78 |
|