Update app.py

app.py

@@ -5,6 +5,12 @@ from decord import VideoReader, cpu
 from PIL import Image
 from torchvision.transforms.functional import InterpolationMode
 from transformers import AutoModel, AutoTokenizer
+from fastapi import FastAPI, UploadFile, File
+from typing import List
+from io import BytesIO
+
+# FastAPI app initialization
+app = FastAPI()
 
 # Device Configuration
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -22,21 +28,6 @@ def build_transform(input_size):
     ])
     return transform
 
-def find_closest_aspect_ratio(aspect_ratio, target_ratios, width, height, image_size):
-    best_ratio_diff = float('inf')
-    best_ratio = (1, 1)
-    area = width * height
-    for ratio in target_ratios:
-        target_aspect_ratio = ratio[0] / ratio[1]
-        ratio_diff = abs(aspect_ratio - target_aspect_ratio)
-        if ratio_diff < best_ratio_diff:
-            best_ratio_diff = ratio_diff
-            best_ratio = ratio
-        elif ratio_diff == best_ratio_diff:
-            if area > 0.5 * image_size * image_size * ratio[0] * ratio[1]:
-                best_ratio = ratio
-    return best_ratio
-
 def dynamic_preprocess(image, min_num=1, max_num=12, image_size=448, use_thumbnail=False):
     orig_width, orig_height = image.size
     aspect_ratio = orig_width / orig_height
@@ -46,16 +37,11 @@ def dynamic_preprocess(image, min_num=1, max_num=12, image_size=448, use_thumbna
         i * j <= max_num and i * j >= min_num)
     target_ratios = sorted(target_ratios, key=lambda x: x[0] * x[1])
 
-    target_aspect_ratio = find_closest_aspect_ratio(
-        aspect_ratio, target_ratios, orig_width, orig_height, image_size)
-
-    target_width = image_size * target_aspect_ratio[0]
-    target_height = image_size * target_aspect_ratio[1]
-    blocks = target_aspect_ratio[0] * target_aspect_ratio[1]
-
+    target_width = image_size * target_ratios[0][0]
+    target_height = image_size * target_ratios[0][1]
     resized_img = image.resize((target_width, target_height))
     processed_images = []
-    for i in range(blocks):
+    for i in range(target_ratios[0][0] * target_ratios[0][1]):
         box = (
             (i % (target_width // image_size)) * image_size,
             (i // (target_width // image_size)) * image_size,
@@ -64,13 +50,12 @@ def dynamic_preprocess(image, min_num=1, max_num=12, image_size=448, use_thumbna
         )
         split_img = resized_img.crop(box)
         processed_images.append(split_img)
-    assert len(processed_images) == blocks
     if use_thumbnail and len(processed_images) != 1:
         thumbnail_img = image.resize((image_size, image_size))
         processed_images.append(thumbnail_img)
     return processed_images
 
-def load_image(image_file, input_size=448, max_num=12):
+def load_image(image_file: BytesIO, input_size=448, max_num=12):
     image = Image.open(image_file).convert('RGB')
     transform = build_transform(input_size=input_size)
     images = dynamic_preprocess(image, image_size=input_size, use_thumbnail=True, max_num=max_num)
@@ -78,38 +63,6 @@ def load_image(image_file, input_size=448, max_num=12):
     pixel_values = torch.stack(pixel_values).to(device)
     return pixel_values
 
-def load_video(video_path, bound=None, input_size=448, max_num=1, num_segments=32):
-    vr = VideoReader(video_path, ctx=cpu(0), num_threads=1)
-    max_frame = len(vr) - 1
-    fps = float(vr.get_avg_fps())
-
-    pixel_values_list, num_patches_list = [], []
-    transform = build_transform(input_size=input_size)
-    frame_indices = get_index(bound, fps, max_frame, first_idx=0, num_segments=num_segments)
-    for frame_index in frame_indices:
-        img = Image.fromarray(vr[frame_index].asnumpy()).convert('RGB')
-        img = dynamic_preprocess(img, image_size=input_size, use_thumbnail=True, max_num=max_num)
-        pixel_values = [transform(tile) for tile in img]
-        pixel_values = torch.stack(pixel_values)
-        num_patches_list.append(pixel_values.shape[0])
-        pixel_values_list.append(pixel_values)
-    pixel_values = torch.cat(pixel_values_list)
-    return pixel_values, num_patches_list
-
-def get_index(bound, fps, max_frame, first_idx=0, num_segments=32):
-    if bound:
-        start, end = bound[0], bound[1]
-    else:
-        start, end = -100000, 100000
-    start_idx = max(first_idx, round(start * fps))
-    end_idx = min(round(end * fps), max_frame)
-    seg_size = float(end_idx - start_idx) / num_segments
-    frame_indices = np.array([
-        int(start_idx + (seg_size / 2) + np.round(seg_size * idx))
-        for idx in range(num_segments)
-    ])
-    return frame_indices
-
 # Load Model
 path = 'OpenGVLab/InternVL2_5-1B'
 model = AutoModel.from_pretrained(
@@ -119,3 +72,14 @@ model = AutoModel.from_pretrained(
     trust_remote_code=True
 ).eval().to(device)
 tokenizer = AutoTokenizer.from_pretrained(path, trust_remote_code=True, use_fast=False)
+
+@app.post("/predict")
+async def predict(file: UploadFile = File(...), question: str = "Describe the image"):
+    # Load and preprocess the image
+    file_bytes = BytesIO(await file.read())
+    pixel_values = load_image(file_bytes)
+    
+    # Generate a response
+    generation_config = dict(max_new_tokens=1024, do_sample=True)
+    response, _ = model.chat(tokenizer, pixel_values, question, generation_config)
+    return {"question": question, "response": response}