From b02926df1393df311db734af149fb9faf4389cbe Mon Sep 17 00:00:00 2001
From: Bruno Seoane <brunoseoaneamarillo@gmail.com>
Date: Sat, 22 Oct 2022 20:24:04 -0300
Subject: [PATCH] Moved moodels to their own file and extracted base64
 conversion to its own function

---
 modules/api/api.py    | 17 ++++++-----------
 modules/api/models.py |  8 ++++++++
 2 files changed, 14 insertions(+), 11 deletions(-)
 create mode 100644 modules/api/models.py

diff --git a/modules/api/api.py b/modules/api/api.py
index a5136b4b..c17d7580 100644
--- a/modules/api/api.py
+++ b/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):
-    images: list[str] = Field(default=None, title="Image", description="The generated image in base64 format.")
-    parameters: Json
-    info: Json
+def img_to_base64(img):
+    buffer = io.BytesIO()
+    img.save(buffer, format="png")
+    return base64.b64encode(buffer.getvalue())
 
 class Api:
     def __init__(self, app, queue_lock):
@@ -41,15 +41,10 @@ class Api:
         with self.queue_lock:
             processed = process_images(p)
         
-        b64images = []
-        for i in processed.images:
-            buffer = io.BytesIO()
-            i.save(buffer, format="png")
-            b64images.append(base64.b64encode(buffer.getvalue()))
+        b64images = list(map(img_to_base64, processed.images))
 
         return TextToImageResponse(images=b64images, parameters=json.dumps(vars(txt2imgreq)), info=json.dumps(processed.info))
         
-        
 
     def img2imgapi(self):
         raise NotImplementedError
diff --git a/modules/api/models.py b/modules/api/models.py
new file mode 100644
index 00000000..a7d247d8
--- /dev/null
+++ b/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
+
+    
\ No newline at end of file