Upload 4 files

Dockerfile

@@ -0,0 +1,22 @@
+FROM pytorch/pytorch:2.3.1-cuda12.1-cudnn8-runtime
+
+RUN apt-get update && apt-get install -y wget
+RUN useradd -m -u 1000 user
+
+USER user
+WORKDIR /app
+
+ENV PATH="/home/user/.local/bin:$PATH"
+ENV TRANSFORMERS_CACHE=/home/user/.cache/huggingface
+ENV TORCH_CUDA_ARCH_LIST="8.0+PTX"
+
+RUN wget https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.0.4/flash_attn-2.7.3+cu121torch2.3-cp310-cp310-linux_x86_64.whl
+RUN pip install ./flash_attn-2.7.3+cu121torch2.3-cp310-cp310-linux_x86_64.whl && rm flash_attn-2.7.3+cu121torch2.3-cp310-cp310-linux_x86_64.whl
+
+COPY --chown=user requirements.txt .
+RUN pip install --upgrade pip setuptools wheel
+RUN pip install --no-cache-dir -r requirements.txt
+
+COPY --chown=user . .
+
+CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "7860"]

endpoint_handler.py

@@ -0,0 +1,91 @@
+import torch
+from PIL import Image
+from transformers import AutoModel, AutoTokenizer
+from io import BytesIO
+import base64
+from huggingface_hub import login
+from huggingface_hub import login
+import os
+
+class EndpointHandler:
+    def __init__(self, model_dir=None):
+        print("[Init] Initializing EndpointHandler...")
+        self.load_model()
+
+    def load_model(self):
+        hf_token = os.getenv("HF_TOKEN")
+        model_path = "openbmb/MiniCPM-o-2_6"  # use model repo name directly
+
+        if hf_token:
+            print("[Auth] Logging into Hugging Face Hub with token...")
+            login(token=hf_token)
+
+        print(f"[Model Load] Loading model from: {model_path}")
+        try:
+            self.tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+            self.model = AutoModel.from_pretrained(
+                model_path,
+                trust_remote_code=True,
+                attn_implementation='sdpa',
+                torch_dtype='auto',  # safer on Spaces
+                init_vision=True,
+                init_audio=False,
+                init_tts=False
+            ).eval().cuda()
+            print("[Model Load] Model successfully loaded and moved to CUDA.")
+        except Exception as e:
+            print(f"[Model Load Error] {e}")
+            raise RuntimeError(f"Failed to load model: {e}")
+        
+    def load_image(self, image_base64):
+        try:
+            print("[Image Load] Decoding base64 image...")
+            image_bytes = base64.b64decode(image_base64)
+            image = Image.open(BytesIO(image_bytes)).convert("RGB")
+            print("[Image Load] Image successfully decoded and converted to RGB.")
+            return image
+        except Exception as e:
+            print(f"[Image Load Error] {e}")
+            raise ValueError(f"Failed to open image from base64 string: {e}")
+
+    def predict(self, request):
+        print(f"[Predict] Received request: {request}")
+
+        image_base64 = request.get("inputs", {}).get("image")
+        question = request.get("inputs", {}).get("question")
+        stream = request.get("inputs", {}).get("stream", False)
+
+        if not image_base64 or not question:
+            print("[Predict Error] Missing 'image' or 'question' in the request.")
+            return {"error": "Missing 'image' or 'question' in inputs."}
+
+        try:
+            image = self.load_image(image_base64)
+            msgs = [{"role": "user", "content": [image, question]}]
+
+            print(f"[Predict] Asking model with question: {question}")
+            print("[Predict] Starting chat inference...")
+
+            res = self.model.chat(
+                image=None,
+                msgs=msgs,
+                tokenizer=self.tokenizer,
+                sampling=True,
+                stream=stream
+            )
+
+            if stream:
+                for new_text in res:
+                    yield {"output": new_text}
+            else:
+                generated_text = "".join(res)
+                print("[Predict] Inference complete.")
+                return {"output": generated_text}
+
+        except Exception as e:
+            print(f"[Predict Error] {e}")
+            return {"error": str(e)}
+
+    def __call__(self, data):
+        print("[__call__] Invoked handler with data.")
+        return self.predict(data)

main.py

@@ -0,0 +1,90 @@
+from fastapi import FastAPI
+from fastapi.responses import JSONResponse, StreamingResponse
+from pydantic import BaseModel
+import types
+import json
+from pydantic import validator
+from endpoint_handler import EndpointHandler  # your handler file
+import base64
+
+app = FastAPI()
+
+handler = None
+
+@app.on_event("startup")
+async def load_handler():
+    global handler
+    handler = EndpointHandler()
+
+class PredictInput(BaseModel):
+    image: str       # base64-encoded image string
+    question: str
+    stream: bool = False
+
+    @validator("question")
+    def question_not_empty(cls, v):
+        if not v.strip():
+            raise ValueError("Question must not be empty")
+        return v
+
+    @validator("image")
+    def valid_base64_and_size(cls, v):
+        try:
+            decoded = base64.b64decode(v, validate=True)
+        except Exception:
+            raise ValueError("`image` must be valid base64")
+        if len(decoded) > 10 * 1024 * 1024:  # 10 MB limit
+            raise ValueError("Image exceeds 10 MB after decoding")
+        return v
+
+class PredictRequest(BaseModel):
+    inputs: PredictInput
+
+@app.get("/")
+async def root():
+    return {"message": "FastAPI app is running on Hugging Face"}
+
+@app.post("/predict")
+async def predict_endpoint(payload: PredictRequest):
+    """
+    Handles prediction requests by processing the input payload and returning the prediction result.
+    Args:
+        payload (PredictRequest): The request payload containing the input data for prediction, including image, question, and stream flag.
+    Returns:
+        JSONResponse: If a ValueError occurs, returns a JSON response with an error message and status code 400.
+        JSONResponse: If any other exception occurs, returns a JSON response with a generic error message and status code 500.
+        StreamingResponse: If the prediction result is a generator (streaming), returns a streaming response with event-stream media type, yielding prediction chunks as JSON.
+    Notes:
+        - Logs the received question for debugging purposes.
+        - Handles both standard and streaming prediction results.
+        - Structured JSON messages are sent to indicate the end of the stream or errors during streaming.
+    """
+    print(f"[Request] Received question: {payload.inputs.question}")
+    
+    data = {
+        "inputs": {
+            "image": payload.inputs.image,
+            "question": payload.inputs.question,
+            "stream": payload.inputs.stream
+        }
+    }
+    
+    try:
+        result = handler.predict(data)
+    except ValueError as ve:
+        return JSONResponse({"error": str(ve)}, status_code=400)
+    except Exception as e:
+        return JSONResponse({"error": "Internal server error"}, status_code=500)
+
+    if isinstance(result, types.GeneratorType):
+        def event_stream():
+            try:
+                for chunk in result:
+                    yield f"data: {json.dumps(chunk)}\n\n"
+                # Return structured JSON to indicate end of stream
+                yield f"data: {json.dumps({'end': True})}\n\n"
+            except Exception as e:
+                # Return structured JSON to indicate error
+                yield f"data: {json.dumps({'error': str(e)})}\n\n"
+    return StreamingResponse(event_stream(), media_type="text/event-stream")
+

requirements.txt

@@ -0,0 +1,19 @@
+Pillow==10.1.0
+torch==2.3.1
+torchaudio==2.3.1
+torchvision==0.18.1
+transformers==4.44.2
+librosa==0.9.0
+soundfile==0.12.1
+vector-quantize-pytorch==1.18.5
+vocos==0.1.0
+decord
+moviepy
+einops
+accelerate
+openbmb 
+fastapi
+uvicorn[standard]
+timm>=0.6.13
+sentencepiece>=0.1.99
+python-multipart