Skip to content

Commit 75b67ee

Browse files
authored
Fix bare base64 not accept
1 parent 828438b commit 75b67ee

File tree

1 file changed

+10
-3
lines changed

1 file changed

+10
-3
lines changed

modules/api/api.py

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
import time
44
import uvicorn
55
from threading import Lock
6+
from io import BytesIO
67
from gradio.processing_utils import encode_pil_to_base64, decode_base64_to_file, decode_base64_to_image
78
from fastapi import APIRouter, Depends, FastAPI, HTTPException
89
from fastapi.security import HTTPBasic, HTTPBasicCredentials
@@ -13,7 +14,7 @@
1314
from modules.api.models import *
1415
from modules.processing import StableDiffusionProcessingTxt2Img, StableDiffusionProcessingImg2Img, process_images
1516
from modules.extras import run_extras, run_pnginfo
16-
from PIL import PngImagePlugin
17+
from PIL import PngImagePlugin,Image
1718
from modules.sd_models import checkpoints_list
1819
from modules.realesrgan_model import get_realesrgan_models
1920
from typing import List
@@ -133,7 +134,10 @@ def img2imgapi(self, img2imgreq: StableDiffusionImg2ImgProcessingAPI):
133134

134135
mask = img2imgreq.mask
135136
if mask:
136-
mask = decode_base64_to_image(mask)
137+
if mask.startswith("data:image/"):
138+
mask = decode_base64_to_image(mask)
139+
else:
140+
mask = Image.open(BytesIO(base64.b64decode(mask)))
137141

138142
populate = img2imgreq.copy(update={ # Override __init__ params
139143
"sd_model": shared.sd_model,
@@ -147,7 +151,10 @@ def img2imgapi(self, img2imgreq: StableDiffusionImg2ImgProcessingAPI):
147151

148152
imgs = []
149153
for img in init_images:
150-
img = decode_base64_to_image(img)
154+
if img.startswith("data:image/"):
155+
img = decode_base64_to_image(img)
156+
else:
157+
img = Image.open(BytesIO(base64.b64decode(img)))
151158
imgs = [img] * p.batch_size
152159

153160
p.init_images = imgs

0 commit comments

Comments
 (0)