Spaces:
Runtime error
Runtime error
Commit
·
e500465
1
Parent(s):
1d1a8c3
added flash-attn to requirements.txt
Browse files- app.py +7 -1
- 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(
|
|
|
|
|
|
|
|
|
|
|
|
|
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
|