xzerus commited on
Commit
42699b5
·
verified ·
1 Parent(s): 7d64545

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +3 -4
app.py CHANGED
@@ -6,11 +6,11 @@ import gradio as gr
6
  # Load the model
7
  model = AutoModel.from_pretrained(
8
  'OpenGVLab/InternVL2_5-1B',
9
- torch_dtype=torch.bfloat16,
10
  low_cpu_mem_usage=True,
11
  trust_remote_code=True,
12
  use_flash_attn=False # Disable Flash Attention
13
- ).cuda().eval()
14
 
15
  # Load the image processor
16
  image_processor = CLIPImageProcessor.from_pretrained('OpenGVLab/InternVL2_5-1B')
@@ -23,9 +23,8 @@ def process_image(image):
23
 
24
  # Preprocess the image
25
  pixel_values = image_processor(images=image, return_tensors='pt').pixel_values
26
- pixel_values = pixel_values.to(torch.bfloat16).cuda()
27
 
28
- # Run the model
29
  outputs = model(pixel_values)
30
 
31
  # Assuming the model returns embeddings or features
 
6
  # Load the model
7
  model = AutoModel.from_pretrained(
8
  'OpenGVLab/InternVL2_5-1B',
9
+ torch_dtype=torch.float32, # Use float32 for CPU compatibility
10
  low_cpu_mem_usage=True,
11
  trust_remote_code=True,
12
  use_flash_attn=False # Disable Flash Attention
13
+ ).eval() # Do not move to CUDA, leave on CPU
14
 
15
  # Load the image processor
16
  image_processor = CLIPImageProcessor.from_pretrained('OpenGVLab/InternVL2_5-1B')
 
23
 
24
  # Preprocess the image
25
  pixel_values = image_processor(images=image, return_tensors='pt').pixel_values
 
26
 
27
+ # Run the model on CPU
28
  outputs = model(pixel_values)
29
 
30
  # Assuming the model returns embeddings or features