Update TextGen/router.py

TextGen/router.py

@@ -2,12 +2,13 @@ import os
 import time
 from langchain_core.pydantic_v1 import BaseModel, Field
 from fastapi import FastAPI, HTTPException, Query, Request
-from fastapi.responses import StreamingResponse
+from fastapi.responses import StreamingResponse,Response
 from fastapi.middleware.cors import CORSMiddleware
 
 from langchain.chains import LLMChain
 from langchain.prompts import PromptTemplate
 from TextGen.suno import custom_generate_audio, get_audio_information,generate_lyrics
+from TextGen.diffusion import generate_image
 from coqui import predict
 from langchain_google_genai import (
     ChatGoogleGenerativeAI,
@@ -229,4 +230,9 @@ async def generate_song():
        
     })
     infos=get_audio_information(f"{data[0]['id']},{data[1]['id']}")
-    return infos
+    return infos
+
+@app.post('/generate_image')
+def Imagen(prompt):
+    image_bytes=generate_image(prompt)
+    return Response(content=image_bytes, media_type="image/png")