Updated main.py

main.py

@@ -1,22 +1,47 @@
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
 from transformers import AutoModelForCausalLM, AutoProcessor
 from PIL import Image
 import requests
+import torch
 
-# Load the processor and model
-processor = AutoProcessor.from_pretrained('allenai/Molmo-7B-D-0924', trust_remote_code=True, device_map='auto')
-model = AutoModelForCausalLM.from_pretrained('allenai/Molmo-7B-D-0924', trust_remote_code=True, device_map='auto')
+# Define the FastAPI app
+app = FastAPI()
 
-# Download an image
-image_url = "https://picsum.photos/id/237/536/354"
-image = Image.open(requests.get(image_url, stream=True).raw)
+# Initialize model and processor at startup
+processor = AutoProcessor.from_pretrained('allenai/Molmo-7B-D-0924', trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained('allenai/Molmo-7B-D-0924', trust_remote_code=True)
 
-# Process the image with some input text
-inputs = processor(images=[image], text="Describe this image.")
-inputs = {k: v.to(model.device).unsqueeze(0) for k, v in inputs.items()}
+# Move the model to GPU if available
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
 
-# Generate text based on the input
-output = model.generate_from_batch(inputs, max_new_tokens=200)
+# Request body structure
+class GenerateRequest(BaseModel):
+    image_url: str
+    text_input: str
 
-# Decode and print the generated text
-generated_text = processor.tokenizer.decode(output[0], skip_special_tokens=True)
-print(generated_text)
+# API root endpoint
+@app.get("/")
+def root():
+    return {"message": "Molmo-7B-D API is up and running!"}
+
+# Text generation endpoint
+@app.post("/generate/")
+def generate_text(request: GenerateRequest):
+    try:
+        # Fetch image from URL
+        response = requests.get(request.image_url, stream=True)
+        image = Image.open(response.raw)
+
+        # Preprocess inputs
+        inputs = processor(images=[image], text=request.text_input, return_tensors="pt").to(device)
+
+        # Generate text
+        output_ids = model.generate(inputs["input_ids"], max_new_tokens=200)
+        generated_text = processor.tokenizer.decode(output_ids[0], skip_special_tokens=True)
+
+        return {"generated_text": generated_text}
+
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))