Spaces:
Running
on
Zero
Running
on
Zero
Update app.py
Browse files
app.py
CHANGED
@@ -6,8 +6,11 @@ import torch
|
|
6 |
import spaces # <-- NEW: import spaces for ZeroGPU
|
7 |
import gradio as gr
|
8 |
from PIL import Image
|
9 |
-
from transformers import AutoProcessor, AutoModelForVision2Seq
|
10 |
from peft import PeftModel
|
|
|
|
|
|
|
11 |
|
12 |
# 0. Environment tweaks for Accelerate (unchanged)
|
13 |
os.environ["ACCELERATE_USE_SLOW_RETRIEVAL"] = "true"
|
@@ -22,9 +25,11 @@ dtype = torch.float16 # <-- use float16 for GPU
|
|
22 |
# 2. Load processor
|
23 |
processor = AutoProcessor.from_pretrained(BASE_MODEL, trust_remote_code=True)
|
24 |
|
|
|
25 |
# 3. Load base model on CPU; stream shards to save RAM
|
26 |
base = AutoModelForVision2Seq.from_pretrained(
|
27 |
BASE_MODEL,
|
|
|
28 |
torch_dtype=dtype,
|
29 |
low_cpu_mem_usage=True,
|
30 |
device_map={"": "cpu"},
|
|
|
6 |
import spaces # <-- NEW: import spaces for ZeroGPU
|
7 |
import gradio as gr
|
8 |
from PIL import Image
|
9 |
+
from transformers import AutoProcessor, AutoModelForVision2Seq, AutoConfig
|
10 |
from peft import PeftModel
|
11 |
+
import transformers, accelerate, torch, numpy as np
|
12 |
+
|
13 |
+
print("VERSIONS:", transformers.__version__, accelerate.__version__, torch.__version__, np.__version__)
|
14 |
|
15 |
# 0. Environment tweaks for Accelerate (unchanged)
|
16 |
os.environ["ACCELERATE_USE_SLOW_RETRIEVAL"] = "true"
|
|
|
25 |
# 2. Load processor
|
26 |
processor = AutoProcessor.from_pretrained(BASE_MODEL, trust_remote_code=True)
|
27 |
|
28 |
+
cfg = AutoConfig.from_pretrained(BASE_MODEL, trust_remote_code=True)
|
29 |
# 3. Load base model on CPU; stream shards to save RAM
|
30 |
base = AutoModelForVision2Seq.from_pretrained(
|
31 |
BASE_MODEL,
|
32 |
+
config=cfg,
|
33 |
torch_dtype=dtype,
|
34 |
low_cpu_mem_usage=True,
|
35 |
device_map={"": "cpu"},
|