ColdSlim commited on
Commit
15f54ff
·
verified ·
1 Parent(s): 603802f

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +6 -1
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"},