Update TextGen/router.py

TextGen/router.py

@@ -53,7 +53,8 @@ class SongRequest(BaseModel):
 class Message(BaseModel):
     npc: str | None  = None
     messages: List[str] | None = None
-    
+class ImageGen(BaseModel):
+    prompt: str | None  = None  
 class VoiceMessage(BaseModel):
     npc: str | None  = None
     input: str | None = None
@@ -233,6 +234,6 @@ async def generate_song():
     return infos
 
 @app.post('/generate_image')
-def Imagen(prompt):
-    image_bytes=generate_image(prompt)
+def Imagen(prompt:ImageGen=None):
+    image_bytes=generate_image(ImageGen.prompt)
     return Response(content=image_bytes, media_type="image/png")