sprakhil commited on
Commit
f86a1bd
·
1 Parent(s): 32e3531

adding AutoModelForImageToText

Browse files
Files changed (1) hide show
  1. app.py +3 -2
app.py CHANGED
@@ -1,7 +1,7 @@
1
  import streamlit as st
2
  from PIL import Image
3
  import torch
4
- from transformers import Qwen2VLForConditionalGeneration, AutoProcessor, pipeline
5
  from colpali_engine.models import ColPali, ColPaliProcessor
6
  import os
7
 
@@ -9,7 +9,8 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
9
 
10
  hf_token = os.getenv('HF_TOKEN')
11
  try:
12
- model = pipeline("image-to-text", model="google/paligemma-3b-mix-448", use_auth_token=hf_token)
 
13
  except Exception as e:
14
  st.error(f"Error loading image-to-text model: {e}")
15
  st.stop()
 
1
  import streamlit as st
2
  from PIL import Image
3
  import torch
4
+ from transformers import Qwen2VLForConditionalGeneration, AutoProcessor, AutoModelForImageToText
5
  from colpali_engine.models import ColPali, ColPaliProcessor
6
  import os
7
 
 
9
 
10
  hf_token = os.getenv('HF_TOKEN')
11
  try:
12
+ processor = AutoProcessor.from_pretrained("google/paligemma-3b-mix-448", use_auth_token=hf_token)
13
+ model = AutoModelForImageToText.from_pretrained("google/paligemma-3b-mix-448", use_auth_token=hf_token)
14
  except Exception as e:
15
  st.error(f"Error loading image-to-text model: {e}")
16
  st.stop()