Vishal1122 commited on
Commit
53275a6
·
verified ·
1 Parent(s): a78a54f

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +1 -54
app.py CHANGED
@@ -106,60 +106,7 @@ def id_extractor(image: Image.Image) -> Dict:
106
  generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
107
 
108
  )
109
- resp = output_text[-1].replace("```json", "").replace("```", "").strip()
110
- return json.loads(resp)
111
-
112
-
113
- import torch
114
- from PIL import Image
115
- from transformers import AutoProcessor, AutoModelForVision2Seq
116
- from transformers.image_utils import load_image
117
-
118
- DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
119
-
120
-
121
-
122
- def id_details(image: Image.Image) -> Dict:
123
-
124
- if image is None:
125
- # Return empty dictionary and make the output invisible
126
- return {}, gr.update(visible=False)
127
-
128
-
129
- # Initialize processor and model
130
- processor = AutoProcessor.from_pretrained("HuggingFaceTB/SmolVLM-500M-Instruct")
131
- model = AutoModelForVision2Seq.from_pretrained(
132
- "HuggingFaceTB/SmolVLM-500M-Instruct",
133
- torch_dtype=torch.bfloat16,
134
- _attn_implementation="flash_attention_2" if DEVICE == "cuda" else "eager",
135
- ).to(DEVICE)
136
-
137
- # Create input messages
138
- messages = [
139
- {
140
- "role": "user",
141
- "content": [
142
- {"type": "image"},
143
- {"type": "text", "text": "Can you describe this image?"}
144
- ]
145
- },
146
- ]
147
-
148
- # Prepare inputs
149
- prompt = processor.apply_chat_template(messages, add_generation_prompt=True)
150
- inputs = processor(text=prompt, images=[image], return_tensors="pt")
151
- inputs = inputs.to(DEVICE)
152
-
153
- # Generate outputs
154
- generated_ids = model.generate(**inputs, max_new_tokens=500)
155
- generated_texts = processor.batch_decode(
156
- generated_ids,
157
- skip_special_tokens=True,
158
- )
159
- return generated_texts
160
-
161
-
162
-
163
 
164
  # Define the Gradio interface for the ID extractor
165
  id_interface = gr.Interface(
 
106
  generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
107
 
108
  )
109
+ return output_text
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
110
 
111
  # Define the Gradio interface for the ID extractor
112
  id_interface = gr.Interface(