53 lines
1.7 KiB
Python
53 lines
1.7 KiB
Python
import os
|
|
|
|
from pyrogram import Client, filters
|
|
from pyrogram.types import Message
|
|
from utils.misc import modules_help, prefix
|
|
from utils.scripts import format_exc, import_library
|
|
|
|
clarifai = import_library('clarifai')
|
|
|
|
from clarifai.client.model import Model # noqa: E402
|
|
|
|
|
|
@Client.on_message(filters.command('cdxl', prefix) & filters.me)
|
|
async def cdxl(c: Client, message: Message):
|
|
try:
|
|
chat_id = message.chat.id
|
|
await message.edit_text('<code>Please Wait...</code>')
|
|
|
|
if len(message.command) > 1:
|
|
prompt = message.text.split(maxsplit=1)[1]
|
|
elif message.reply_to_message:
|
|
prompt = message.reply_to_message.text
|
|
else:
|
|
await message.edit_text(f'<b>Usage: </b><code>{prefix}vdxl [prompt/reply to prompt]</code>')
|
|
return
|
|
|
|
inference_params = {'width': 1024, 'height': 1024, 'steps': 50, 'cfg_scale': 9.0}
|
|
|
|
model_prediction = Model(
|
|
'https://clarifai.com/stability-ai/stable-diffusion-2/models/stable-diffusion-xl'
|
|
).predict_by_bytes(prompt.encode(), input_type='text', inference_params=inference_params)
|
|
|
|
output_base64 = model_prediction.outputs[0].data.image.base64
|
|
|
|
with open('sdxl_out.png', 'wb') as f:
|
|
f.write(output_base64)
|
|
|
|
await message.delete()
|
|
await c.send_photo(
|
|
chat_id,
|
|
photo='sdxl_out.png',
|
|
caption=f'<b>Prompt:</b><code>{prompt}</code>',
|
|
)
|
|
os.remove('sdxl_out.png')
|
|
|
|
except Exception as e:
|
|
await message.edit_text(f'An error occurred: {format_exc(e)}')
|
|
|
|
|
|
modules_help['cdxl'] = {
|
|
'cdxl [prompt/reply to prompt]*': 'Text to Image with SDXL model',
|
|
}
|