added image endpoint

main.py

@@ -1,10 +1,9 @@
 # main.py
-from fastapi import FastAPI,status
+from fastapi import FastAPI,status,Response
 from fastapi.responses import StreamingResponse
-
-from models import load_text_model,generate_text,load_audio_model,generate_audio
+from models import load_text_model,generate_text,load_audio_model,generate_audio,load_image_model, generate_image
 from schemas import VoicePresets
-from utils import audio_array_to_buffer
+from utils import audio_array_to_buffer,img_to_bytes
 app = FastAPI() 
 
 @app.get("/")
@@ -31,4 +30,13 @@ def serve_text_to_audio_model_controller(
     output, sample_rate = generate_audio(processor, model, prompt, preset)
     return StreamingResponse(
         audio_array_to_buffer(output, sample_rate), media_type="audio/wav"
-    )
+    )
+
+
+@app.get("/generate/image",
+         responses={status.HTTP_200_OK: {"content": {"image/png": {}}}}, 
+         response_class=Response) 
+def serve_text_to_image_model_controller(prompt: str):
+    pipe = load_image_model()
+    output = generate_image(pipe, prompt) 
+    return Response(content=img_to_bytes(output), media_type="image/png") 

models.py

@@ -3,6 +3,8 @@
 import torch
 from transformers import Pipeline, pipeline,AutoProcessor, AutoModel, BarkProcessor, BarkModel
 from schemas import VoicePresets
+from diffusers import DiffusionPipeline, StableDiffusionInpaintPipelineLegacy
+from PIL import Image
 import numpy as np
 
 prompt = "How to set up a FastAPI project?"
@@ -60,4 +62,18 @@ def generate_text(pipe: Pipeline, prompt: str, temperature: float = 0.7) -> str:
         top_p=0.95,
     ) 
     output = predictions[0]["generated_text"].split("</s>\n<|assistant|>\n")[-1] 
+    return output
+
+
+def load_image_model() -> StableDiffusionInpaintPipelineLegacy:
+    pipe = DiffusionPipeline.from_pretrained(
+        "segmind/tiny-sd", torch_dtype=torch.float32,
+        device=device
+    ) 
+    return pipe
+
+def generate_image(
+    pipe: StableDiffusionInpaintPipelineLegacy, prompt: str
+) -> Image.Image:
+    output = pipe(prompt, num_inference_steps=10).images[0] 
     return output

requirements.txt

@@ -4,4 +4,5 @@ transformers
 torch
 pydantic
 bitsandbytes
-soundfile
+soundfile
+diffusers

utils.py

@@ -1,9 +1,18 @@
 from io import BytesIO
 import soundfile
 import numpy as np
+from PIL import Image
 
 def audio_array_to_buffer(audio_array: np.array, sample_rate: int) -> BytesIO:
     buffer = BytesIO()
     soundfile.write(buffer, audio_array, sample_rate, format="wav") 
     buffer.seek(0)
-    return buffer
+    return buffer
+
+
+def img_to_bytes(
+    image: Image.Image, img_format: Literal["PNG", "JPEG"] = "PNG"
+) -> bytes:
+    buffer = BytesIO()
+    image.save(buffer, format=img_format)
+    return buffer.getvalue()