Improve preview in image generation of Gemini

Fix issue in Bing Image Decorator
This commit is contained in:
Heiner Lohaus 2024-02-11 09:26:02 +01:00
parent 228392eaa2
commit 9aeae65b9b
3 changed files with 8 additions and 5 deletions

View File

@ -87,7 +87,7 @@ class CreateImagesProvider(BaseProvider):
messages.insert(0, {"role": "system", "content": self.system_message})
buffer = ""
for chunk in self.provider.create_completion(model, messages, stream, **kwargs):
if buffer or "<" in chunk:
if isinstance(chunk, str) and buffer or "<" in chunk:
buffer += chunk
if ">" in buffer:
match = re.search(r'<img data-prompt="(.*?)">', buffer)

View File

@ -142,13 +142,15 @@ class Gemini(AsyncGeneratorProvider):
if image_prompt:
images = [image[0][3][3] for image in response_part[4][0][12][7][0]]
resolved_images = []
preview = []
for image in images:
async with session.get(image, allow_redirects=False) as fetch:
image = fetch.headers["location"]
async with session.get(image, allow_redirects=False) as fetch:
image = fetch.headers["location"]
resolved_images.append(image)
yield ImageResponse(resolved_images, image_prompt, {"orginal_links": images})
preview.append(image.replace('=s512', '=s200'))
yield ImageResponse(resolved_images, image_prompt, {"orginal_links": images, "preview": preview})
def build_request(
prompt: str,

View File

@ -187,7 +187,7 @@ def to_base64_jpg(image: Image, compression_rate: float) -> str:
image.save(output_buffer, format="JPEG", quality=int(compression_rate * 100))
return base64.b64encode(output_buffer.getvalue()).decode()
def format_images_markdown(images, alt: str, preview: str = None) -> str:
def format_images_markdown(images: Union[str, list], alt: str, preview: Union[str, list] = None) -> str:
"""
Formats the given images as a markdown string.
@ -202,9 +202,10 @@ def format_images_markdown(images, alt: str, preview: str = None) -> str:
if isinstance(images, str):
images = f"[![{alt}]({preview.replace('{image}', images) if preview else images})]({images})"
else:
if not isinstance(preview, list):
preview = [preview.replace('{image}', image) if preview else image for image in images]
images = [
f"[![#{idx+1} {alt}]({preview.replace('{image}', image) if preview else image})]({image})"
for idx, image in enumerate(images)
f"[![#{idx+1} {alt}]({preview[idx]})]({image})" for idx, image in enumerate(images)
]
images = "\n".join(images)
start_flag = "<!-- generated images start -->\n"