Commit b02926df authored by Bruno Seoane's avatar Bruno Seoane

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

parent 1b4d0473
...@@ -4,17 +4,17 @@ from modules.sd_samplers import all_samplers ...@@ -4,17 +4,17 @@ from modules.sd_samplers import all_samplers
import modules.shared as shared import modules.shared as shared
import uvicorn import uvicorn
from fastapi import APIRouter, HTTPException from fastapi import APIRouter, HTTPException
from pydantic import BaseModel, Field, Json
import json import json
import io import io
import base64 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) 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: class Api:
def __init__(self, app, queue_lock): def __init__(self, app, queue_lock):
...@@ -41,16 +41,11 @@ class Api: ...@@ -41,16 +41,11 @@ class Api:
with self.queue_lock: with self.queue_lock:
processed = process_images(p) 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)) return TextToImageResponse(images=b64images, parameters=json.dumps(vars(txt2imgreq)), info=json.dumps(processed.info))
def img2imgapi(self): def img2imgapi(self):
raise NotImplementedError raise NotImplementedError
......
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
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment