Jaykintecblic commited on
Commit
9d87924
1 Parent(s): 7eb7583

Update handler.py

Browse files
Files changed (1) hide show
  1. handler.py +29 -34
handler.py CHANGED
@@ -1,21 +1,23 @@
1
- from typing import Dict, Any
2
- from fastapi import FastAPI, File, UploadFile
3
- from fastapi.responses import StreamingResponse
4
  from PIL import Image
5
  import torch
6
  from transformers import AutoModelForCausalLM, AutoProcessor
7
  from transformers.image_utils import to_numpy_array, PILImageResampling, ChannelDimension
8
  from transformers.image_transforms import resize, to_channel_dimension_format
9
- import json
10
- import io
11
-
12
- app = FastAPI()
13
 
14
  class EndpointHandler:
15
  def __init__(self, model_path: str):
16
  self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
17
- self.processor = AutoProcessor.from_pretrained(model_path)
18
- self.model = AutoModelForCausalLM.from_pretrained(model_path, trust_remote_code=True, torch_dtype=torch.bfloat16).to(self.device)
 
 
 
 
 
 
 
 
19
  self.image_seq_len = self.model.config.perceiver_config.resampler_n_latents
20
  self.bos_token = self.processor.tokenizer.bos_token
21
  self.bad_words_ids = self.processor.tokenizer(["<image>", "<fake_token_around_image>"], add_special_tokens=False).input_ids
@@ -42,32 +44,25 @@ class EndpointHandler:
42
  image = to_channel_dimension_format(image, ChannelDimension.FIRST)
43
  return torch.tensor(image)
44
 
45
- async def generate_responses(self, image: Image.Image):
46
- try:
47
- inputs = self.processor.tokenizer(
48
- f"{self.bos_token}<fake_token_around_image>{'<image>' * self.image_seq_len}<fake_token_around_image>",
49
- return_tensors="pt",
50
- add_special_tokens=False,
51
- )
52
- inputs["pixel_values"] = self.processor.image_processor([image], transform=self.custom_transform)
53
- inputs = {k: v.to(self.device) for k, v in inputs.items()}
54
 
55
- generated_ids = self.model.generate(**inputs, bad_words_ids=self.bad_words_ids, max_length=2048, early_stopping=True)
56
- generated_text = self.processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
57
- yield json.dumps({"label": generated_text, "score": 1.0}) + '\n'
58
-
59
- except torch.cuda.CudaError as e:
60
- yield json.dumps({"error": f"CUDA error: {e}"}) + '\n'
61
- except Exception as e:
62
- yield json.dumps({"error": f"Unexpected error: {e}"}) + '\n'
63
 
64
- handler = EndpointHandler(model_path="Jaykintecblic/Html")
 
 
 
 
 
 
65
 
66
- @app.post("/")
67
- async def handle_request(file: UploadFile = File(...)):
68
- image = Image.open(io.BytesIO(await file.read()))
69
- return StreamingResponse(handler.generate_responses(image), media_type="application/json")
 
 
70
 
71
- if __name__ == "__main__":
72
- import uvicorn
73
- uvicorn.run(app, host="0.0.0.0", port=8080)
 
1
+ from typing import Dict, List, Any
 
 
2
  from PIL import Image
3
  import torch
4
  from transformers import AutoModelForCausalLM, AutoProcessor
5
  from transformers.image_utils import to_numpy_array, PILImageResampling, ChannelDimension
6
  from transformers.image_transforms import resize, to_channel_dimension_format
 
 
 
 
7
 
8
  class EndpointHandler:
9
  def __init__(self, model_path: str):
10
  self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
11
+ self.processor = AutoProcessor.from_pretrained(
12
+ model_path,
13
+ # token=api_token
14
+ )
15
+ self.model = AutoModelForCausalLM.from_pretrained(
16
+ model_path,
17
+ # token=api_token,
18
+ trust_remote_code=True,
19
+ torch_dtype=torch.bfloat16,
20
+ ).to(self.device)
21
  self.image_seq_len = self.model.config.perceiver_config.resampler_n_latents
22
  self.bos_token = self.processor.tokenizer.bos_token
23
  self.bad_words_ids = self.processor.tokenizer(["<image>", "<fake_token_around_image>"], add_special_tokens=False).input_ids
 
44
  image = to_channel_dimension_format(image, ChannelDimension.FIRST)
45
  return torch.tensor(image)
46
 
47
+ def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
48
+ image = data.get("inputs")
 
 
 
 
 
 
 
49
 
50
+ if isinstance(image, str):
51
+ image = Image.open(image)
 
 
 
 
 
 
52
 
53
+ inputs = self.processor.tokenizer(
54
+ f"{self.bos_token}<fake_token_around_image>{'<image>' * self.image_seq_len}<fake_token_around_image>",
55
+ return_tensors="pt",
56
+ add_special_tokens=False,
57
+ )
58
+ inputs["pixel_values"] = self.processor.image_processor([image], transform=self.custom_transform)
59
+ inputs = {k: v.to(self.device) for k, v in inputs.items()}
60
 
61
+ generated_ids = self.model.generate(**inputs, bad_words_ids=self.bad_words_ids, max_length=2048, early_stopping=True)
62
+ generated_text = self.processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
63
+ # print(generated_text)
64
+ # return {"text": generated_text}
65
+ # Format the output as an array of dictionaries with 'label' and 'score'
66
+ output = [{"label": generated_text, "score": 1.0}]
67
 
68
+ return output