lyttt commited on
Commit
0e2f2fd
·
1 Parent(s): ab2ce79

Add application file

Browse files
Files changed (1) hide show
  1. app.py +3 -1
app.py CHANGED
@@ -7,7 +7,9 @@ model_name_or_path = "lyttt/VLV_captioner"
7
  model = AutoModel.from_pretrained(model_name_or_path, revision="master", trust_remote_code=True,low_cpu_mem_usage=False)
8
 
9
  def greet(image):
10
- image = image.convert("RGB")
 
 
11
  with torch.no_grad():
12
  outputs = model([image]).generated_text[0]
13
  return output
 
7
  model = AutoModel.from_pretrained(model_name_or_path, revision="master", trust_remote_code=True,low_cpu_mem_usage=False)
8
 
9
  def greet(image):
10
+ if image.dtype != np.uint8:
11
+ image = (np.clip(image, 0, 1) * 255).astype(np.uint8)
12
+ image = Image.fromarray(image, mode='RGB')
13
  with torch.no_grad():
14
  outputs = model([image]).generated_text[0]
15
  return output