Update TextGen/router.py

TextGen/router.py

@@ -10,7 +10,7 @@ from langchain.chains import LLMChain
 from langchain.prompts import PromptTemplate
 from TextGen.suno import custom_generate_audio, get_audio_information,generate_lyrics
 from TextGen.gemini import generate_story,place_objects,generate_map_markdown
-#from TextGen.diffusion import generate_image
+from TextGen.diffusion import generate_image
 #from coqui import predict
 from langchain_google_genai import (
     ChatGoogleGenerativeAI,
@@ -56,7 +56,9 @@ class Message(BaseModel):
     npc: str | None  = None
     messages: List[str] | None = None
 class ImageGen(BaseModel):
-    prompt: str | None  = None  
+    name:str
+    prompt: str
+    
 class VoiceMessage(BaseModel):
     npc: str | None  = None
     input: str | None = None
@@ -321,15 +323,15 @@ async def generate_song():
     infos=get_audio_information(f"{data[0]['id']},{data[1]['id']}")
     return infos
 
-#@app.post('/generate_image')
-#def Imagen(image:ImageGen=None):
-#    pil_image =generate_image(image.prompt)
-#    
-#
-#    # Convert the PIL Image to bytes
-#    img_byte_arr = BytesIO()
-#    pil_image.save(img_byte_arr, format='PNG')
-#    img_byte_arr = img_byte_arr.getvalue()
+@app.post('/generate_image')
+def Imagen(image:ImageGen=None):
+    pil_image =generate_image(image.prompt)
+    
+
+    # Convert the PIL Image to bytes
+    img_byte_arr = BytesIO()
+    pil_image.save(img_byte_arr, format='PNG')
+    img_byte_arr = img_byte_arr.getvalue()
 #
     # Return the image as a PNG response
-#    return Response(content=img_byte_arr, media_type="image/png")
+    return Response(content=img_byte_arr, media_type="image/png")