Update app.py

app.py

@@ -8,6 +8,7 @@ import uvicorn
 import soundfile as sf
 import imageio
 from typing import Dict, Optional
+import torch  # Import torch
 
 AWS_ACCESS_KEY_ID = os.getenv("AWS_ACCESS_KEY_ID")
 AWS_SECRET_ACCESS_KEY = os.getenv("AWS_SECRET_ACCESS_KEY")
@@ -137,7 +138,7 @@ async def generate(request: GenerateRequest, model_resources: tuple = Depends(ge
         new_continuation_id = continuation_id if continuation_id else os.urandom(16).hex()
         active_generations[new_continuation_id] = {"model_name": model_name, "output": generated_text}
 
-        return JSONResponse({"text": generated_text, "continuation_id": new_continuation_id})
+        return JSONResponse({"text": generated_text, "continuation_id": new_continuation_id, "model_name": model_name})
 
     except HTTPException as http_err:
         raise http_err
@@ -187,7 +188,7 @@ async def generate_image(request: GenerateRequest):
         image = image_generator(request.input_text)[0]
         new_continuation_id = os.urandom(16).hex()
         active_generations[new_continuation_id] = {"model_name": request.model_name, "output": "Image generated successfully"}
-        return JSONResponse({"url": "Image generated successfully", "continuation_id": new_continuation_id})
+        return JSONResponse({"url": "Image generated successfully", "continuation_id": new_continuation_id, "model_name": request.model_name})
 
     except HTTPException as http_err:
         raise http_err
@@ -204,7 +205,7 @@ async def generate_text_to_speech(request: GenerateRequest):
         output = tts_pipeline(request.input_text)
         new_continuation_id = os.urandom(16).hex()
         active_generations[new_continuation_id] = {"model_name": request.model_name, "output": "Audio generated successfully"}
-        return JSONResponse({"url": "Audio generated successfully", "continuation_id": new_continuation_id})
+        return JSONResponse({"url": "Audio generated successfully", "continuation_id": new_continuation_id, "model_name": request.model_name})
 
     except HTTPException as http_err:
         raise http_err
@@ -221,7 +222,7 @@ async def generate_video(request: GenerateRequest):
         output = video_pipeline(request.input_text)
         new_continuation_id = os.urandom(16).hex()
         active_generations[new_continuation_id] = {"model_name": request.model_name, "output": "Video generated successfully"}
-        return JSONResponse({"url": "Video generated successfully", "continuation_id": new_continuation_id})
+        return JSONResponse({"url": "Video generated successfully", "continuation_id": new_continuation_id, "model_name": request.model_name})
 
     except HTTPException as http_err:
         raise http_err