patrickluan11 commited on
Commit
5131c02
路
verified 路
1 Parent(s): 8f2d1b6

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +7 -1
app.py CHANGED
@@ -1,6 +1,11 @@
 
1
  import tempfile
2
  from pathlib import Path
3
 
 
 
 
 
4
  import gradio as gr
5
  import spaces
6
  import torch
@@ -20,7 +25,8 @@ def get_pipeline():
20
  MODEL_ID,
21
  torch_dtype=torch.bfloat16,
22
  )
23
- PIPE.to("cuda")
 
24
  print("Model loaded on CUDA", flush=True)
25
  return PIPE
26
 
 
1
+ import os
2
  import tempfile
3
  from pathlib import Path
4
 
5
+ # Alguns workers ZeroGPU n茫o exp玫em NVML completamente ao allocator do PyTorch.
6
+ # Desativar o cache evita o assert interno sem alterar o resultado do modelo.
7
+ os.environ.setdefault("PYTORCH_NO_CUDA_MEMORY_CACHING", "1")
8
+
9
  import gradio as gr
10
  import spaces
11
  import torch
 
25
  MODEL_ID,
26
  torch_dtype=torch.bfloat16,
27
  )
28
+ # O offload evita mover o pipeline inteiro de uma s贸 vez para a GPU.
29
+ PIPE.enable_model_cpu_offload()
30
  print("Model loaded on CUDA", flush=True)
31
  return PIPE
32