lewtun HF staff commited on
Commit
1f70fa4
·
verified ·
1 Parent(s): d433054

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +3 -3
app.py CHANGED
@@ -1,5 +1,5 @@
1
  import gradio as gr
2
- from transformers import LlavaNextProcessor, LlavaNextForConditionalGeneration, TextIteratorStreamer
3
  from threading import Thread
4
  import re
5
  import time
@@ -13,9 +13,9 @@ login(token=os.environ["HF_TOKEN"])
13
 
14
  MODEL_ID = os.environ.get("MODEL_ID", "llava-hf/llava-v1.6-mistral-7b-hf")
15
  REVISION = os.environ.get("MODEL_REVISION", "main")
16
- processor = LlavaNextProcessor.from_pretrained(MODEL_ID, revision=REVISION)
17
 
18
- model = LlavaNextForConditionalGeneration.from_pretrained(MODEL_ID, revision=REVISION, torch_dtype=torch.float16, low_cpu_mem_usage=True)
19
  model.to("cuda:0")
20
 
21
  @spaces.GPU
 
1
  import gradio as gr
2
+ from transformers import LlavaProcessor, LlavaForConditionalGeneration, TextIteratorStreamer
3
  from threading import Thread
4
  import re
5
  import time
 
13
 
14
  MODEL_ID = os.environ.get("MODEL_ID", "llava-hf/llava-v1.6-mistral-7b-hf")
15
  REVISION = os.environ.get("MODEL_REVISION", "main")
16
+ processor = LlavaProcessor.from_pretrained(MODEL_ID, revision=REVISION)
17
 
18
+ model = LlavaForConditionalGeneration.from_pretrained(MODEL_ID, revision=REVISION, torch_dtype=torch.float16, low_cpu_mem_usage=True)
19
  model.to("cuda:0")
20
 
21
  @spaces.GPU