Fix create_async in CreateImagesProvider

This commit is contained in:
Heiner Lohaus 2024-01-10 20:08:06 +01:00
parent 9aa76da0fa
commit 4343082d37

View File

@ -71,14 +71,17 @@ class CreateImagesProvider(BaseProvider):
) -> str:
messages.insert(0, {"role": "system", "content": self.system_message})
response = await self.provider.create_async(model, messages, **kwargs)
matches = re.findall(r'(<img data-prompt="(.*?)">)', result)
matches = re.findall(r'(<img data-prompt="(.*?)">)', response)
results = []
for _, prompt in matches:
placeholders = []
for placeholder, prompt in matches:
if placeholder not in placeholders:
results.append(self.create_images_async(prompt))
placeholders.append(placeholder)
results = await asyncio.gather(*results)
for idx, result in enumerate(results):
placeholder = matches[idx][0]
placeholder = placeholder[idx]
if self.include_placeholder:
result = placeholder + result
response = response.replace(placeholder, result)
return result
return response