Update app.py

app.py

@@ -2,57 +2,64 @@ import gradio as gr
 import torch
 from PIL import Image
 from transformers import AutoProcessor, LlavaForConditionalGeneration
-import time
-import os
+from io import BytesIO
 import requests
 import json
+import time
+
+# Load processor and model
+processor = AutoProcessor.from_pretrained("llava-hf/llava-1.5-7b-hf")
 
-# Load model
-processor = AutoProcessor.from_pretrained("llava-hf/llava-1.5-7b-hf")  # Using 7B to fit in free tier
 model = LlavaForConditionalGeneration.from_pretrained(
     "llava-hf/llava-1.5-7b-hf",
     torch_dtype=torch.float16,
-    low_cpu_mem_usage=True
-).to("cuda")
+    device_map="auto"
+)
 
-# API function
-def api_endpoint(request):
+# Core inference function
+def generate_response(user_message, system_prompt=None, image=None, max_tokens=1024, temperature=0.7):
+    if system_prompt:
+        prompt = f"<image>\n{system_prompt}\n{user_message}"
+    else:
+        prompt = f"<image>\n{user_message}"
+
+    inputs = processor(prompt, image, return_tensors="pt").to(model.device)
+
+    with torch.inference_mode():
+        output = model.generate(
+            **inputs,
+            max_new_tokens=max_tokens,
+            do_sample=True,
+            temperature=temperature,
+        )
+
+    response_text = processor.decode(output[0], skip_special_tokens=True)
+    return response_text
+
+# API-style function for programmatic access
+def api_endpoint(request: gr.Request):
     try:
-        # Parse request
-        data = json.loads(request)
+        data = request.json
         user_message = data.get("user_message", "")
         system_prompt = data.get("system_prompt", None)
         image_url = data.get("image_url", None)
         max_tokens = data.get("max_tokens", 1024)
         temperature = data.get("temperature", 0.7)
-        
-        # Process image if provided
+
         image_data = None
         if image_url:
-            response = requests.get(image_url)
-            image_data = Image.open(BytesIO(response.content)).convert("RGB")
-        
-        # Prepare prompt
-        if system_prompt:
-            prompt = f"<s>[INST] {system_prompt} [/INST]\n{user_message}"
-        else:
-            prompt = user_message
-            
-        # Generate response
-        inputs = processor(prompt, image_data, return_tensors="pt").to(model.device)
-        
-        with torch.inference_mode():
-            output = model.generate(
-                **inputs,
-                max_new_tokens=max_tokens,
-                do_sample=True,
-                temperature=temperature,
-            )
-        
-        response_text = processor.decode(output[0], skip_special_tokens=True)
-        
-        # Return response
-        return json.dumps({
+            image_response = requests.get(image_url)
+            image_data = Image.open(BytesIO(image_response.content)).convert("RGB")
+
+        response_text = generate_response(
+            user_message=user_message,
+            system_prompt=system_prompt,
+            image=image_data,
+            max_tokens=max_tokens,
+            temperature=temperature
+        )
+
+        return gr.Response(json.dumps({
             "id": f"chatcmpl-{int(time.time())}",
             "object": "chat.completion",
             "created": int(time.time()),
@@ -65,16 +72,17 @@ def api_endpoint(request):
                 "index": 0,
                 "finish_reason": "stop"
             }]
-        })
-        
+        }), media_type="application/json")
+
     except Exception as e:
-        return json.dumps({"error": str(e)})
+        return gr.Response(json.dumps({"error": str(e)}), media_type="application/json")
+
 
-# Gradio Interface
+# Gradio UI
 with gr.Blocks() as demo:
-    gr.Markdown("# LLaVA API Demo")
-    
-    with gr.Tab("API Test UI"):
+    gr.Markdown("# 🔍 LLaVA API Demo")
+
+    with gr.Tab("Test UI"):
         with gr.Row():
             with gr.Column():
                 user_message = gr.Textbox(label="User Message", lines=3)
@@ -83,123 +91,20 @@ with gr.Blocks() as demo:
                 max_tokens = gr.Slider(label="Max Tokens", minimum=1, maximum=2048, value=1024, step=1)
                 temperature = gr.Slider(label="Temperature", minimum=0.1, maximum=2.0, value=0.7, step=0.1)
                 submit_btn = gr.Button("Generate Response")
-            
             with gr.Column():
                 output = gr.Textbox(label="Response", lines=10)
-    
-    with gr.Tab("API Documentation"):
-        gr.Markdown("""
-        ## API Endpoint Documentation
-        
-        **URL**: `https://YOUR-USERNAME-llava-api.hf.space/api/`
-        
-        **Method**: POST
-        
-        **Request Body**:
-        ```json
-        {
-            "user_message": "Describe this image",
-            "system_prompt": "You are a helpful assistant",
-            "image_url": "https://example.com/image.jpg",
-            "max_tokens": 1024,
-            "temperature": 0.7
-        }
-        ```
-        
-        **Response**:
-        ```json
-        {
-            "id": "chatcmpl-1234567890",
-            "object": "chat.completion",
-            "created": 1683123456,
-            "model": "llava-1.5-7b",
-            "choices": [
-                {
-                    "message": {
-                        "role": "assistant",
-                        "content": "Response text here"
-                    },
-                    "index": 0,
-                    "finish_reason": "stop"
-                }
-            ]
-        }
-        ```
-        
-        **Example Python Client**:
-        ```python
-        import requests
-        import json
-        
-        def query_llava_api(api_url, user_message, system_prompt=None, image_url=None):
-            payload = {
-                "user_message": user_message,
-                "max_tokens": 1024
-            }
-            
-            if system_prompt:
-                payload["system_prompt"] = system_prompt
-                
-            if image_url:
-                payload["image_url"] = image_url
-                
-            response = requests.post(api_url, json=payload)
-            return response.json()
-            
-        # Example usage
-        result = query_llava_api(
-            "https://YOUR-USERNAME-llava-api.hf.space/api/",
-            "What's in this image?",
-            image_url="https://example.com/image.jpg"
-        )
-        print(result["choices"][0]["message"]["content"])
-        ```
-        """)
-    
-    # API endpoint
-    gr.Interface(
-        fn=api_endpoint,
-        inputs=gr.Textbox(),
-        outputs=gr.Textbox(),
-        api_name="api"
-    )
-    
-    # Connect UI to function
-    def process_inputs(message, system, img, tokens, temp):
-        # Create payload
-        payload = {
-            "user_message": message,
-            "max_tokens": tokens,
-            "temperature": temp
-        }
-        
-        if system:
-            payload["system_prompt"] = system
-            
-        # Process image
-        if img is not None:
-            # For demo purposes, we use the image directly
-            # In a real API, you'd need to handle image URLs
-            inputs = processor(message, img, return_tensors="pt").to(model.device)
-            
-            with torch.inference_mode():
-                output = model.generate(
-                    **inputs,
-                    max_new_tokens=tokens,
-                    do_sample=True,
-                    temperature=temp,
-                )
-            
-            response_text = processor.decode(output[0], skip_special_tokens=True)
-            return response_text
-        
-        # If no image, process text only
-        return api_endpoint(json.dumps(payload))
-    
+
+    def on_submit(message, system, image, tokens, temp):
+        return generate_response(message, system, image, tokens, temp)
+
     submit_btn.click(
-        process_inputs,
+        fn=on_submit,
         inputs=[user_message, system_prompt, image_input, max_tokens, temperature],
         outputs=output
     )
 
+# API endpoint
+demo.api("/api")(api_endpoint)
+
+# Launch
 demo.launch()