kls/sd_telegram
1
0
Fork 0
forked from soaska/sd_telegram
Code Pull Requests Actions Packages Activity

just_gen - instant image generation for channels and chats

Browse Source 
channel_post_handler - catch a post from a channel and generate an instant image
999999999 an attempt to limit the sending of large images (not working yet)
This commit is contained in:
Mikhail Guseletov 2024-01-18 12:54:14 +07:00
parent 48e92b085e
commit aa245f93a0
1 changed files with 45 additions and 16 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

61
bot.py
View File

@ -31,6 +31,7 @@ from translate import Translator
import base64 import base64
from pathlib import Path from pathlib import Path
import logging import logging
import sys
import vk_api import vk_api
from vk_api import VkUpload #https://github.com/python273/vk_api from vk_api import VkUpload #https://github.com/python273/vk_api
from ok_api import OkApi, Upload # https://github.com/needkirem/ok_api from ok_api import OkApi, Upload # https://github.com/needkirem/ok_api
@ -120,6 +121,7 @@ dataParams = {"img_thumb": "true",
"sd_model_checkpoint": "", "sd_model_checkpoint": "",
"use_prompt": "true", "use_prompt": "true",
"json_prompt": "false", "json_prompt": "false",
"just_gen": "false",
"send_vk": "false"} "send_vk": "false"}
dataOld = data.copy() dataOld = data.copy()
dataOldParams = dataParams.copy() dataOldParams = dataParams.copy()
@ -174,6 +176,8 @@ def pilToImages(res, typeImages="tg"):
image_buffer = io.BytesIO() image_buffer = io.BytesIO()
image.save(image_buffer, format="PNG") image.save(image_buffer, format="PNG")
image_buffer.seek(0) image_buffer.seek(0)
if image_buffer.getbuffer().nbytes > 999999999 and typeImages == 'tg':
typeImages = 'thumbs'
# картинка в телеге # картинка в телеге
if typeImages == "tg": if typeImages == "tg":
media_group.append(types.InputMediaPhoto(media=image_buffer, caption=seed)) media_group.append(types.InputMediaPhoto(media=image_buffer, caption=seed))
@ -528,7 +532,7 @@ async def rnd_script(message, typeScript):
text=dataPromptOld text=dataPromptOld
) )
for i, number in enumerate(numbers): for i, number in enumerate(numbers):
time.sleep(5) time.sleep(1)
for itemTxt in data['prompt'].split(';'): for itemTxt in data['prompt'].split(';'):
if typeScript == 'models': if typeScript == 'models':
api.util_wait_for_ready() api.util_wait_for_ready()
@ -564,15 +568,15 @@ async def rnd_script(message, typeScript):
# show thumb/tg/real # show thumb/tg/real
async def show_thumbs(chat_id, res): async def show_thumbs(chat_id, res):
if dataParams["img_thumb"] == "true" or dataParams["img_thumb"] == "True": if str(data['img_thumb']).lower() == 'true':
await bot.send_media_group( await bot.send_media_group(
chat_id=chat_id, media=pilToImages(res, "thumbs") chat_id=chat_id, media=pilToImages(res, "thumbs")
) )
if dataParams["img_tg"] == "true" or dataParams["img_tg"] == "True": if str(data['img_tg']).lower() == 'true':
await bot.send_media_group( await bot.send_media_group(
chat_id=chat_id, media=pilToImages(res, "tg") chat_id=chat_id, media=pilToImages(res, "tg")
) )
if dataParams["img_real"] == "true" or dataParams["img_real"] == "True": if str(data['img_real']).lower() == 'true':
messages = await bot.send_media_group( messages = await bot.send_media_group(
chat_id=chat_id, chat_id=chat_id,
media=pilToImages(res, "real") media=pilToImages(res, "real")
@ -861,13 +865,15 @@ async def inl_fp(message: Union[types.Message, types.CallbackQuery]) -> None:
"use_prompt": "true", "use_prompt": "true",
"json_prompt": "false"} "json_prompt": "false"}
if m == 'fp_sdxl': if m == 'fp_sdxl':
data['steps'] = 5 data['enable_hr'] = 'True'
data['denoising_strength'] = '0.3'
data['steps'] = 15
data['sampler_name'] = 'DPM++ SDE Karras' data['sampler_name'] = 'DPM++ SDE Karras'
data['enable_hr'] = 'false'
data['cfg_scale'] = '4' data['cfg_scale'] = '4'
data['width'] = '1024' data['width'] = '1024'
data['height'] = '1024' data['height'] = '1024'
data['restore_faces'] = 'false' data['restore_faces'] = 'false'
data['hr_upscaler'] = '4x_NMKD-Siax_200k'
data['do_not_save_grid'] = 'true' data['do_not_save_grid'] = 'true'
data['negative_prompt'] = 'FastNegativeV2' data['negative_prompt'] = 'FastNegativeV2'
data['save_images'] = 'true' data['save_images'] = 'true'
@ -876,7 +882,8 @@ async def inl_fp(message: Union[types.Message, types.CallbackQuery]) -> None:
"img_real": "true", "img_real": "true",
"stop_sd": "true", "stop_sd": "true",
"use_prompt": "true", "use_prompt": "true",
"json_prompt": "false"} "json_prompt": "false",
"just_gen": "false"}
txt = f"JSON отредактирован\n{getJson()}\n{getJson(1)}" txt = f"JSON отредактирован\n{getJson()}\n{getJson(1)}"
await getKeyboardUnion(txt, message, keyboard, '') await getKeyboardUnion(txt, message, keyboard, '')
@ -925,15 +932,15 @@ async def inl_gen(message: Union[types.Message, types.CallbackQuery]) -> None:
# TODO try catch if wrong data # TODO try catch if wrong data
res = await api.txt2img(**data) res = await api.txt2img(**data)
# show_thumbs dont work because use_async # show_thumbs dont work because use_async
if dataParams["img_thumb"] == "true" or dataParams["img_thumb"] == "True": if str(data['img_thumb']).lower() == 'true':
await bot.send_media_group( await bot.send_media_group(
chat_id=chatId, media=pilToImages(res, "thumbs") chat_id=chatId, media=pilToImages(res, "thumbs")
) )
if dataParams["img_tg"] == "true" or dataParams["img_tg"] == "True": if str(data['img_tg']).lower() == 'true':
await bot.send_media_group( await bot.send_media_group(
chat_id=chatId, media=pilToImages(res, "tg") chat_id=chatId, media=pilToImages(res, "tg")
) )
if dataParams["img_real"] == "true" or dataParams["img_real"] == "True": if str(data['img_real']).lower() == 'true':
messages = await bot.send_media_group( messages = await bot.send_media_group(
chat_id=chatId, chat_id=chatId,
media=pilToImages(res, "real") media=pilToImages(res, "real")
@ -1341,10 +1348,24 @@ async def inl_yes_no(callback: types.CallbackQuery) -> None:
reply_markup=keyboard, reply_markup=keyboard,
) )
# отлов поста с канала и мгновенная генерация если включен just_gen
@dp.channel_post_handler()
async def handle_channel_post(message: types.Message):
logging.info("handle_channel_post")
chatId = message.chat.id
data["prompt"] = message.text
if str(dataParams['just_gen']).lower() == 'true':
data["use_async"] = "True"
res = await api.txt2img(**data)
await bot.send_media_group(
chat_id=chatId, media=pilToImages(res, "tg")
)
# Ввели любой текст # Ввели любой текст
@dp.message_handler(lambda message: True) @dp.message_handler(lambda message: True)
async def change_json(message: types.Message): async def change_json(message: types.Message):
logging.info("change_json") logging.info("change_json")
chatId = message.chat.id
keyboard = InlineKeyboardMarkup(inline_keyboard=[getSet(0), getOpt(0), getStart(0)]) keyboard = InlineKeyboardMarkup(inline_keyboard=[getSet(0), getOpt(0), getStart(0)])
text = message.text text = message.text
nam = text.split()[0][1:] # txt из /txt 321 nam = text.split()[0][1:] # txt из /txt 321
@ -1391,12 +1412,20 @@ async def change_json(message: types.Message):
f"JSON параметры:\n{getJson()}\n{getJson(1)}", reply_markup=keyboard f"JSON параметры:\n{getJson()}\n{getJson(1)}", reply_markup=keyboard
) )
else: else:
data["prompt"] = message.text#translateRuToEng(message.text) data["prompt"] = message.text # translateRuToEng(message.text)
await bot.delete_message(chat_id=message.chat.id, message_id=message.message_id) if str(dataParams['just_gen']).lower() == 'true':
await message.answer( data["use_async"] = "True"
f"Записали промпт. JSON параметры:\n{getJson()}\n{getJson(1)}", res = await api.txt2img(**data)
reply_markup=keyboard, # TODO img_real and data
) await bot.send_media_group(
chat_id=chatId, media=pilToImages(res, "tg")
)
else:
await bot.delete_message(chat_id=message.chat.id, message_id=message.message_id)
await message.answer(
f"Записали промпт. JSON параметры:\n{getJson()}\n{getJson(1)}",
reply_markup=keyboard,
)
# Ввели ответ на change_json # Ввели ответ на change_json
@dp.message_handler(state=Form) @dp.message_handler(state=Form)