merve HF staff commited on
Commit
c4055fa
1 Parent(s): 41ebd25

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +6 -9
app.py CHANGED
@@ -14,7 +14,7 @@ from transformers import AutoProcessor, Blip2ForConditionalGeneration
14
  DESCRIPTION = "# [BLIP-2](https://github.com/salesforce/LAVIS/tree/main/projects/blip2)"
15
 
16
  if not torch.cuda.is_available():
17
- DESCRIPTION += "\n<p>Running on CPU 🥶 This demo does not work on CPU.</p>"
18
 
19
  device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
20
 
@@ -22,15 +22,12 @@ MODEL_ID_OPT_6_7B = "Salesforce/blip2-opt-6.7b"
22
  MODEL_ID_FLAN_T5_XXL = "Salesforce/blip2-flan-t5-xxl"
23
  MODEL_ID = os.getenv("MODEL_ID", MODEL_ID_FLAN_T5_XXL)
24
 
25
- if torch.cuda.is_available():
26
- processor = AutoProcessor.from_pretrained(MODEL_ID)
27
- model = Blip2ForConditionalGeneration.from_pretrained(MODEL_ID, device_map="auto", load_in_8bit=True)
28
- else:
29
- processor = None
30
- model = None
31
 
32
 
33
- @spaces.GPU
34
  def generate_caption(
35
  image: PIL.Image.Image,
36
  decoding_method: str,
@@ -54,7 +51,7 @@ def generate_caption(
54
  return result
55
 
56
 
57
- @spaces.GPU
58
  def answer_question(
59
  image: PIL.Image.Image,
60
  text: str,
 
14
  DESCRIPTION = "# [BLIP-2](https://github.com/salesforce/LAVIS/tree/main/projects/blip2)"
15
 
16
  if not torch.cuda.is_available():
17
+ DESCRIPTION += "\n<p>Running on CPU, might be slow.</p>"
18
 
19
  device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
20
 
 
22
  MODEL_ID_FLAN_T5_XXL = "Salesforce/blip2-flan-t5-xxl"
23
  MODEL_ID = os.getenv("MODEL_ID", MODEL_ID_FLAN_T5_XXL)
24
 
25
+
26
+ processor = AutoProcessor.from_pretrained(MODEL_ID)
27
+ model = Blip2ForConditionalGeneration.from_pretrained(MODEL_ID, device_map="auto", load_in_8bit=True)
28
+
 
 
29
 
30
 
 
31
  def generate_caption(
32
  image: PIL.Image.Image,
33
  decoding_method: str,
 
51
  return result
52
 
53
 
54
+
55
  def answer_question(
56
  image: PIL.Image.Image,
57
  text: str,