Update app.py

app.py

@@ -1,51 +1,79 @@
-from fastapi import FastAPI, HTTPException, Response
+from fastapi import FastAPI, File, UploadFile, HTTPException
+import requests
+import base64
 from pydantic import BaseModel
-from transformers import AutoTokenizer, AutoModel
-import torch
-import numpy as np
-import io
-from scipy.io.wavfile import write
-from PIL import Image
-import riffusion
+from typing import Optional
 
 app = FastAPI()
 
-# Load Riffusion model
-model_name = "riffusion/riffusion-model-v1"
-model = AutoModel.from_pretrained(model_name)
-tokenizer = AutoTokenizer.from_pretrained(model_name)
+# NVIDIA API endpoint and API key
+NVIDIA_API_URL = "https://ai.api.nvidia.com/v1/gr/meta/llama-3.2-90b-vision-instruct/chat/completions"
+API_KEY = "your_nvidia_api_key_here"  # Replace with your actual API key
 
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model.to(device)
+# Request model for text-based input
+class TextRequest(BaseModel):
+    message: str
+    max_tokens: Optional[int] = 512
+    temperature: Optional[float] = 1.0
+    top_p: Optional[float] = 1.0
 
-class MusicRequest(BaseModel):
-    prompt: str
+# Function to call the NVIDIA API
+def call_nvidia_api(payload: dict):
+    headers = {
+        "Authorization": f"Bearer {API_KEY}",
+        "Accept": "application/json",
+    }
+    response = requests.post(NVIDIA_API_URL, headers=headers, json=payload)
+    if response.status_code != 200:
+        raise HTTPException(status_code=response.status_code, detail="NVIDIA API request failed")
+    return response.json()
 
-@app.post("/generate-music/")
-def generate_music(request: MusicRequest):
+# Endpoint for text-based input
+@app.post("/chat/text")
+async def chat_with_text(request: TextRequest):
+    payload = {
+        "model": "meta/llama-3.2-90b-vision-instruct",
+        "messages": [{"role": "user", "content": request.message}],
+        "max_tokens": request.max_tokens,
+        "temperature": request.temperature,
+        "top_p": request.top_p,
+        "stream": False,
+    }
     try:
-        inputs = tokenizer(request.prompt, return_tensors="pt").to(device)
-        with torch.no_grad():
-            spectrogram = model.generate(**inputs)
-        
-        # Convert spectrogram to an image (since Riffusion outputs spectrograms)
-        spectrogram_image = Image.fromarray((spectrogram.cpu().numpy().squeeze() * 255).astype(np.uint8))
-        
-        # Convert spectrogram to audio
-        audio_values, sampling_rate = riffusion.audio_processing.spectrogram_to_audio(spectrogram_image)
-        
-        # Normalize and convert to int16
-        audio_values = np.clip(audio_values * 32767, -32768, 32767).astype(np.int16)
-        
-        # Convert to WAV format
-        audio_bytes = io.BytesIO()
-        write(audio_bytes, sampling_rate, audio_values)
-        audio_bytes.seek(0)
-        
-        return Response(content=audio_bytes.read(), media_type="audio/wav", headers={"Content-Disposition": "attachment; filename=generated_music.wav"})
+        response = call_nvidia_api(payload)
+        return {"response": response["choices"][0]["message"]["content"]}
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
+# Endpoint for image-based input
+@app.post("/chat/image")
+async def chat_with_image(file: UploadFile = File(...)):
+    # Read and encode the image file to base64
+    image_data = await file.read()
+    base64_image = base64.b64encode(image_data).decode("utf-8")
+
+    # Prepare the payload for the NVIDIA API
+    payload = {
+        "model": "meta/llama-3.2-90b-vision-instruct",
+        "messages": [
+            {
+                "role": "user",
+                "content": f'What is in this image? <img src="data:image/png;base64,{base64_image}" />',
+            }
+        ],
+        "max_tokens": 512,
+        "temperature": 1.0,
+        "top_p": 1.0,
+        "stream": False,
+    }
+
+    try:
+        response = call_nvidia_api(payload)
+        return {"response": response["choices"][0]["message"]["content"]}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+# Root endpoint
 @app.get("/")
-def root():
-    return {"message": "Welcome to the Riffusion Music Generation API"}
+async def root():
+    return {"message": "Welcome to the NVIDIA API FastAPI wrapper!"}