Moved moodels to their own file and extracted base64 conversion to its own function

modules/api/api.py

@@ -4,17 +4,17 @@ from modules.sd_samplers import all_samplers
 import modules.shared as shared
 import uvicorn
 from fastapi import APIRouter, HTTPException
-from pydantic import BaseModel, Field, Json
 import json
 import io
 import base64
+from modules.api.models import *
 
 sampler_to_index = lambda name: next(filter(lambda row: name.lower() == row[1].name.lower(), enumerate(all_samplers)), None)
 
-class TextToImageResponse(BaseModel):
+def img_to_base64(img):
-    images: list[str] = Field(default=None, title="Image", description="The generated image in base64 format.")
+    buffer = io.BytesIO()
-    parameters: Json
+    img.save(buffer, format="png")
-    info: Json
+    return base64.b64encode(buffer.getvalue())
 
 class Api:
     def __init__(self, app, queue_lock):
@@ -41,16 +41,11 @@ class Api:
         with self.queue_lock:
             processed = process_images(p)
         
-        b64images = []
+        b64images = list(map(img_to_base64, processed.images))
-        for i in processed.images:
-            buffer = io.BytesIO()
-            i.save(buffer, format="png")
-            b64images.append(base64.b64encode(buffer.getvalue()))
 
         return TextToImageResponse(images=b64images, parameters=json.dumps(vars(txt2imgreq)), info=json.dumps(processed.info))
         
 
-
     def img2imgapi(self):
         raise NotImplementedError
 

modules/api/models.py

@@ -0,0 +1,8 @@
+from pydantic import BaseModel, Field, Json
+
+class TextToImageResponse(BaseModel):
+    images: list[str] = Field(default=None, title="Image", description="The generated image in base64 format.")
+    parameters: Json
+    info: Json
+
+