jiang-cc commited on
Commit
49bfeb1
·
verified ·
1 Parent(s): b0bcd32

fix: load model directly to cuda like ltx-2-TURBO, enable flash_attention_2

Browse files
Files changed (1) hide show
  1. app.py +2 -2
app.py CHANGED
@@ -38,8 +38,8 @@ model = AutoModelForImageTextToText.from_pretrained(
38
  torch_dtype=torch.bfloat16,
39
  attn_implementation=_attn_impl,
40
  trust_remote_code=True,
41
- ).eval()
42
- print(f"[AD-Copilot] Attention: {_attn_impl}", flush=True)
43
 
44
 
45
  # ---------------------------------------------------------------------------
 
38
  torch_dtype=torch.bfloat16,
39
  attn_implementation=_attn_impl,
40
  trust_remote_code=True,
41
+ ).to("cuda").eval()
42
+ print(f"[AD-Copilot] Attention: {_attn_impl} | Device: {model.device}", flush=True)
43
 
44
 
45
  # ---------------------------------------------------------------------------