thankfulcarp commited on
Commit
e500465
·
1 Parent(s): 1d1a8c3

added flash-attn to requirements.txt

Browse files
Files changed (2) hide show
  1. app.py +7 -1
  2. requirements.txt +4 -1
app.py CHANGED
@@ -209,7 +209,13 @@ def enhance_prompt_with_llm(prompt: str):
209
  print("\n🤖 Loading LLM for Prompt Enhancement (first run)...")
210
  try:
211
  # This happens inside the GPU session, so device_map="auto" is correct.
212
- ENHANCER_PIPE_CACHE = pipeline("text-generation", model=ENHANCER_MODEL_ID, torch_dtype=torch.bfloat16, device_map="auto")
 
 
 
 
 
 
213
  print("✅ LLM Prompt Enhancer loaded successfully.")
214
  except Exception as e:
215
  print(f"❌ Error loading LLM enhancer: {e}")
 
209
  print("\n🤖 Loading LLM for Prompt Enhancement (first run)...")
210
  try:
211
  # This happens inside the GPU session, so device_map="auto" is correct.
212
+ ENHANCER_PIPE_CACHE = pipeline(
213
+ "text-generation",
214
+ model=ENHANCER_MODEL_ID,
215
+ torch_dtype=torch.bfloat16,
216
+ device_map="auto",
217
+ model_kwargs={"attn_implementation": "flash_attention_2"}
218
+ )
219
  print("✅ LLM Prompt Enhancer loaded successfully.")
220
  except Exception as e:
221
  print(f"❌ Error loading LLM enhancer: {e}")
requirements.txt CHANGED
@@ -17,4 +17,7 @@ sentencepiece
17
  ftfy
18
  imageio
19
  imageio-ffmpeg
20
- opencv-python
 
 
 
 
17
  ftfy
18
  imageio
19
  imageio-ffmpeg
20
+ opencv-python
21
+
22
+ # Performance
23
+ flash-attn