Update handler.py

handler.py

@@ -1,115 +1,135 @@
+import os
+import json
 import base64
+import tempfile
 from io import BytesIO
 
 import requests
-import torch
 from PIL import Image
-from transformers import (
-    AutoConfig,
-    AutoProcessor,
-    PretrainedConfig,
-    Qwen2_5_VLForConditionalGeneration,
-)
+from swift.llm import PtEngine, RequestConfig, get_model_tokenizer, get_template, InferRequest
+
 
 MODEL_ID = "ghost233lism/GeoAgent"
 
 
 class EndpointHandler:
     def __init__(self, path=""):
-        dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
-
-        # 关键修复：先读 config，再把 dict 子配置转成 PretrainedConfig
-        config = AutoConfig.from_pretrained(
-            MODEL_ID,
-            trust_remote_code=True,
-        )
-
-        if hasattr(config, "text_config") and isinstance(config.text_config, dict):
-            config.text_config = PretrainedConfig.from_dict(config.text_config)
-
-        if hasattr(config, "vision_config") and isinstance(config.vision_config, dict):
-            config.vision_config = PretrainedConfig.from_dict(config.vision_config)
-
-        # 有些模型这里也可能是 dict，一并兜底
-        if hasattr(config, "decoder_config") and isinstance(config.decoder_config, dict):
-            config.decoder_config = PretrainedConfig.from_dict(config.decoder_config)
-
-        self.model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
-            MODEL_ID,
-            config=config,
-            torch_dtype=dtype,
-            device_map="auto",
-            trust_remote_code=True,
-            low_cpu_mem_usage=True,
+        self.model_path = MODEL_ID
+        self.max_new_tokens = 2048
+        self.engine = None
+        self.template = None
+        self._load_model()
+
+    def _load_model(self):
+        model, tokenizer = get_model_tokenizer(self.model_path, model_type="qwen2_5_vl")
+        template_type = model.model_meta.template
+        self.template = get_template(
+            template_type,
+            tokenizer,
+            default_system=self._system_prompt(),
         )
-
-        self.processor = AutoProcessor.from_pretrained(
-            MODEL_ID,
-            trust_remote_code=True,
+        self.engine = PtEngine.from_model_template(
+            model,
+            self.template,
+            max_batch_size=1,
         )
 
-    def _load_image(self, image_ref):
-        if image_ref is None:
-            raise ValueError("Please provide `inputs.image_url` or `inputs.image_base64`.")
-
-        if isinstance(image_ref, str) and image_ref.startswith("http"):
-            resp = requests.get(image_ref, timeout=30)
+    @staticmethod
+    def _system_prompt():
+        return """You are an expert with rich experience in the field of geolocation, skilled at accurately locating the geographic location of images through various clues in the images, such as traffic signs, architectural styles, natural landscapes, etc. At the same time, you are also a mentor in building the chain of thought, able to organize complex ideas into clear and standardized patterns.
+You possess knowledge in multiple disciplines such as geography, cartography, transportation, and architecture, and are able to identify the characteristics of different countries, regions, and locations. At the same time, you have the ability to analyze logic and construct a chain of thought. Task: Output the thought chain and final answer based on the image input by the user. The thought chain includes:
+Country Identification/Regional Guess/Precise Localization.
+
+Possible clues include: National clues: (Example: traffic sign shape/color, language and text, driving direction, architectural style, vegetation and climate characteristics, etc.)
+
+Regional clues: (logo/enterprise, topography, vegetation type, regional traffic signs, dialect/spelling, license plate style, area code/postal code, infrastructure features, etc.)
+Accurate positioning: (road sign text, street name, house number, landmark building, river and lake water system, place attributes such as park/city/commercial district, shop name and storefront, etc.)
+
+Do not output objects that do not exist in the image.
+
+Output strictly in JSON format:
+{
+  "ChainOfThought": {
+    "CountryIdentification": {
+      "Clues": [],
+      "Reasoning": "",
+      "Conclusion": "",
+      "Uncertainty": ""
+    },
+    "RegionalGuess": {
+      "Clues": [],
+      "Reasoning": "",
+      "Conclusion": "",
+      "Uncertainty": ""
+    },
+    "PreciseLocalization": {
+      "Clues": [],
+      "Reasoning": "",
+      "Conclusion": "",
+      "Uncertainty": ""
+    }
+  },
+  "FinalAnswer": "Country; Region; Specific Location"
+}"""
+
+    def _save_temp_image(self, image_url=None, image_base64=None):
+        if image_url:
+            resp = requests.get(image_url, timeout=30)
             resp.raise_for_status()
-            return Image.open(BytesIO(resp.content)).convert("RGB")
-
-        if isinstance(image_ref, str) and image_ref.startswith("data:image"):
-            _, b64data = image_ref.split(",", 1)
-            return Image.open(BytesIO(base64.b64decode(b64data))).convert("RGB")
+            img = Image.open(BytesIO(resp.content)).convert("RGB")
+        elif image_base64:
+            if image_base64.startswith("data:image"):
+                _, image_base64 = image_base64.split(",", 1)
+            img = Image.open(BytesIO(base64.b64decode(image_base64))).convert("RGB")
+        else:
+            raise ValueError("Please provide `inputs.image_url` or `inputs.image_base64`.")
 
-        return Image.open(image_ref).convert("RGB")
+        tmp = tempfile.NamedTemporaryFile(suffix=".jpg", delete=False)
+        img.save(tmp.name, format="JPEG")
+        tmp.close()
+        return tmp.name
 
     def __call__(self, data):
         payload = data.get("inputs", {}) or {}
 
-        prompt = payload.get("prompt", "Please analyze this image and infer its location.")
         image_url = payload.get("image_url")
         image_base64 = payload.get("image_base64")
-        max_new_tokens = int(payload.get("max_new_tokens", 256))
-
-        image = self._load_image(image_url or image_base64)
-
-        messages = [
-            {
-                "role": "user",
-                "content": [
-                    {"type": "image"},
-                    {"type": "text", "text": prompt},
-                ],
-            }
-        ]
-
-        text = self.processor.apply_chat_template(
-            messages,
-            tokenize=False,
-            add_generation_prompt=True,
+        prompt = payload.get(
+            "prompt",
+            "Based on the image, tell me the specific location and your thinking process",
         )
+        max_new_tokens = int(payload.get("max_new_tokens", self.max_new_tokens))
 
-        model_inputs = self.processor(
-            text=[text],
-            images=[image],
-            return_tensors="pt",
-        ).to(self.model.device)
+        image_path = self._save_temp_image(image_url=image_url, image_base64=image_base64)
 
-        with torch.no_grad():
-            output_ids = self.model.generate(
-                **model_inputs,
-                max_new_tokens=max_new_tokens,
+        try:
+            infer_request = InferRequest(
+                messages=[
+                    {
+                        "role": "user",
+                        "content": f"<image>{prompt}",
+                    }
+                ],
+                images=[image_path],
             )
 
-        generated_ids = [
-            out_ids[len(in_ids):]
-            for in_ids, out_ids in zip(model_inputs.input_ids, output_ids)
-        ]
+            request_config = RequestConfig(max_tokens=max_new_tokens)
 
-        output_text = self.processor.batch_decode(
-            generated_ids,
-            skip_special_tokens=True,
-            clean_up_tokenization_spaces=True,
-        )[0]
+            resp = self.engine.infer([infer_request], request_config=request_config)
+            result = resp[0]
 
-        return {"generated_text": output_text}
+            text = ""
+            if hasattr(result, "choices") and result.choices:
+                choice = result.choices[0]
+                if hasattr(choice, "message") and hasattr(choice.message, "content"):
+                    text = choice.message.content
+                elif hasattr(choice, "text"):
+                    text = choice.text
+
+            return {
+                "generated_text": text,
+                "raw_response": str(result),
+            }
+        finally:
+            if os.path.exists(image_path):
+                os.remove(image_path)