fthor commited on
Commit
a76b117
·
1 Parent(s): f192c41

Added back quantization

Browse files
Files changed (1) hide show
  1. app.py +5 -5
app.py CHANGED
@@ -3,17 +3,17 @@ import torch
3
  from transformers import AutoProcessor, LlavaForConditionalGeneration
4
  from transformers import BitsAndBytesConfig
5
 
6
- # quantization_config = BitsAndBytesConfig(
7
- # load_in_4bit=True,
8
- # bnb_4bit_compute_dtype=torch.float16
9
- # )
10
 
11
  model_id = "llava-hf/llava-1.5-7b-hf"
12
 
13
  processor = AutoProcessor.from_pretrained(model_id)
14
  model = LlavaForConditionalGeneration.from_pretrained(
15
  model_id,
16
- # quantization_config=quantization_config,
17
  device_map="auto"
18
  )
19
 
 
3
  from transformers import AutoProcessor, LlavaForConditionalGeneration
4
  from transformers import BitsAndBytesConfig
5
 
6
+ quantization_config = BitsAndBytesConfig(
7
+ load_in_4bit=True,
8
+ bnb_4bit_compute_dtype=torch.float16
9
+ )
10
 
11
  model_id = "llava-hf/llava-1.5-7b-hf"
12
 
13
  processor = AutoProcessor.from_pretrained(model_id)
14
  model = LlavaForConditionalGeneration.from_pretrained(
15
  model_id,
16
+ quantization_config=quantization_config,
17
  device_map="auto"
18
  )
19