vikhyatk commited on
Commit
44f58e0
1 Parent(s): bac7d5d

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +1 -4
app.py CHANGED
@@ -6,9 +6,6 @@ import gradio as gr
6
  from threading import Thread
7
  from transformers import TextIteratorStreamer, AutoTokenizer, AutoModelForCausalLM
8
 
9
- import subprocess
10
- subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
11
-
12
  parser = argparse.ArgumentParser()
13
 
14
  model_id = "vikhyatk/moondream2"
@@ -16,7 +13,7 @@ revision = "2024-04-02"
16
  tokenizer = AutoTokenizer.from_pretrained(model_id, revision=revision)
17
  moondream = AutoModelForCausalLM.from_pretrained(
18
  model_id, trust_remote_code=True, revision=revision,
19
- attn_implementation="flash_attention_2", torch_dtype=torch.float16
20
  )
21
  moondream.eval()
22
 
 
6
  from threading import Thread
7
  from transformers import TextIteratorStreamer, AutoTokenizer, AutoModelForCausalLM
8
 
 
 
 
9
  parser = argparse.ArgumentParser()
10
 
11
  model_id = "vikhyatk/moondream2"
 
13
  tokenizer = AutoTokenizer.from_pretrained(model_id, revision=revision)
14
  moondream = AutoModelForCausalLM.from_pretrained(
15
  model_id, trust_remote_code=True, revision=revision,
16
+ torch_dtype=torch.float32
17
  )
18
  moondream.eval()
19