bgaspra commited on
Commit
bd1b634
·
verified ·
1 Parent(s): 154e0bd

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +11 -8
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 Florence2Model, Florence2Processor
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
- self.model_head = nn.Linear(florence_model.config.hidden_size, num_models)
48
- self.prompt_head = nn.Linear(florence_model.config.hidden_size, vocab_size)
 
49
 
50
- def forward(self, image_features):
51
  # Get Florence embeddings
52
- features = self.florence.get_image_features(**image_features)
 
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=1000):
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 = Florence2Processor.from_pretrained(
68
  "microsoft/Florence-2-large",
69
  trust_remote_code=True
70
  )
71
- self.florence = Florence2Model.from_pretrained(
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