Imaginer/src/provider/stablediffusionlocal.py
0xMRTT 85fc78a222
All checks were successful
ci/woodpecker/push/woodpecker Pipeline was successful
fix: handle json decode error
2023-05-26 20:01:17 +02:00

104 lines
3.4 KiB
Python

from .base import ImaginerProvider
import requests
import socket
import json
from gi.repository import Gtk, Adw, GLib
from PIL import Image, UnidentifiedImageError
import io
import base64
class StableDiffusionLocalProvider(ImaginerProvider):
name = "Local Stable Diffusion"
slug = "stablediffusionlocal"
version = "0.1.0"
url = "https://imaginer.codeberg.page/help/local"
api_url = ""
def __init__(self, win, app, *args, **kwargs):
super().__init__(win, app, *args, **kwargs)
def ask(self, prompt, negative_prompt):
if self.api_url:
try:
payload = json.dumps(
{
"prompt": prompt,
"negative_prompts": negative_prompt if negative_prompt else "",
}
)
headers = {"Content-Type": "application/json"}
url = f"{self.api_url}/sdapi/v1/txt2img"
response = requests.request("POST", url, headers=headers, data=payload)
if response.status_code == 403:
self.no_api_key()
return ""
elif response.status_code != 200:
try:
self.no_api_key(title=response.json()["error"])
return ""
except json.decoder.JSONDecodeError:
self.no_api_key(title="No response")
return ""
response = response.json()
except KeyError:
pass
except socket.gaierror:
self.no_connection()
return ""
except json.decoder.JSONDecodeError:
self.no_api_key(title="No response")
return ""
else:
self.hide_banner()
if response:
try:
img = io.BytesIO(base64.b64decode(response["images"][0]))
return Image.open(img)
except UnidentifiedImageError:
error = json.loads(response)["error"]
self.no_api_key(title=error)
return None
else:
print("No response")
return None
else:
self.no_api_key(title="No API URL selected, you can choose one in preferences")
@property
def require_api_key(self):
return False
def preferences(self, win):
self.pref_win = win
self.expander = Adw.ExpanderRow()
self.expander.props.title = self.name
self.expander.add_action(self.about()) # TODO: in Adw 1.4, use add_suffix
self.expander.add_action(self.enable_switch())
self.api_row = Adw.EntryRow()
self.api_row.connect("apply", self.on_apply)
self.api_row.props.text = self.api_url
self.api_row.props.title = "API Url"
self.api_row.set_show_apply_button(True)
self.api_row.add_suffix(self.how_to_get_a_token())
self.expander.add_row(self.api_row)
return self.expander
def on_apply(self, widget):
self.hide_banner()
self.api_url = self.api_row.get_text()
def save(self):
return {"api_url": self.api_url}
def load(self, data):
if data["api_url"]:
self.api_url = data["api_url"]