hysts HF staff commited on
Commit
321abc5
1 Parent(s): 1a4f7ce
Files changed (1) hide show
  1. app.py +1 -3
app.py CHANGED
@@ -21,13 +21,11 @@ device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
21
  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
 
21
  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
+ assert MODEL_ID in [MODEL_ID_OPT_6_7B, MODEL_ID_FLAN_T5_XXL]
25
 
26
  if torch.cuda.is_available():
27
  processor = AutoProcessor.from_pretrained(MODEL_ID)
28
  model = Blip2ForConditionalGeneration.from_pretrained(MODEL_ID, device_map="auto", load_in_8bit=True)
 
 
 
29
 
30
 
31
  @spaces.GPU