stable-diffusion-telegram-bot/main.py

import json
import requests
import io
import re
import os
import uuid
import base64
from datetime import datetime
from PIL import Image, PngImagePlugin
from pyrogram import Client, filters
from pyrogram.types import InlineKeyboardButton, InlineKeyboardMarkup
from dotenv import load_dotenv

# Load environment variables
load_dotenv()
API_ID = os.environ.get("API_ID", None)
API_HASH = os.environ.get("API_HASH", None)
TOKEN = os.environ.get("TOKEN_givemtxt2img", None)
SD_URL = os.environ.get("SD_URL", None)

app = Client("stable", api_id=API_ID, api_hash=API_HASH, bot_token=TOKEN)
IMAGE_PATH = 'images'  # Do not leave a trailing /

# Ensure IMAGE_PATH directory exists
os.makedirs(IMAGE_PATH, exist_ok=True)

def timestamp():
    return datetime.now().strftime("%Y%m%d-%H%M%S")

def encode_file_to_base64(path):
    with open(path, 'rb') as file:
        return base64.b64encode(file.read()).decode('utf-8')

def decode_and_save_base64(base64_str, save_path):
    with open(save_path, "wb") as file:
        file.write(base64.b64decode(base64_str))

def parse_input(input_string):
    default_payload = {
        "prompt": "",
        "negative_prompt": "ugly, bad face, distorted",
        "enable_hr": False,
        "Sampler": "DPM++ SDE Karras",
        "denoising_strength": 0.35,
        "batch_size": 1,
        "n_iter": 1,
        "steps": 35,
        "cfg_scale": 7,
        "width": 512,
        "height": 512,
        "restore_faces": False,
        "override_settings": {},
        "override_settings_restore_afterwards": True,
    }
    payload = default_payload.copy()
    prompt = []

    matches = re.finditer(r"(\w+):", input_string)
    last_index = 0

    for match in matches:
        key = match.group(1).lower()
        value_start_index = match.end()

        if last_index != match.start():
            prompt.append(input_string[last_index: match.start()].strip())
        last_index = value_start_index
        if key == "ds":
            key = "denoising_strength"
        if key == "ng":
            key = "negative_prompt"
            
        if key in default_payload:
            value_end_index = re.search(r"(?=\s+\w+:|$)", input_string[value_start_index:]).start()
            value = input_string[value_start_index: value_start_index + value_end_index].strip()
            payload[key] = value
            last_index += value_end_index
        else:
            prompt.append(f"{key}:")

    payload["prompt"] = " ".join(prompt)
    if not payload["prompt"]:
        payload["prompt"] = input_string.strip()

    return payload


def call_api(api_endpoint, payload):
    try:
        response = requests.post(f'{SD_URL}/{api_endpoint}', json=payload)
        response.raise_for_status()
        return response.json()
    except requests.RequestException as e:
        print(f"API call failed: {e}")
        return None

def process_images(images, user_id, user_name):
    def generate_unique_name():
        unique_id = str(uuid.uuid4())[:7]
        return f"{user_name}-{unique_id}"

    word = generate_unique_name()

    for i in images:
        image = Image.open(io.BytesIO(base64.b64decode(i.split(",", 1)[0])))

        png_payload = {"image": "data:image/png;base64," + i}
        response2 = requests.post(f"{SD_URL}/sdapi/v1/png-info", json=png_payload)
        response2.raise_for_status()

        pnginfo = PngImagePlugin.PngInfo()
        pnginfo.add_text("parameters", response2.json().get("info"))
        image.save(f"{IMAGE_PATH}/{word}.png", pnginfo=pnginfo)

        return word, response2.json().get("info")

@app.on_message(filters.command(["draw"]))
def draw(client, message):
    msgs = message.text.split(" ", 1)
    if len(msgs) == 1:
        message.reply_text("Format :\n/draw < text to image >\nng: < negative (optional) >\nsteps: < steps value (1-70, optional) >")
        return

    payload = parse_input(msgs[1])
    print(payload)

    K = message.reply_text("Please Wait 10-15 Seconds")
    r = call_api('sdapi/v1/txt2img', payload)

    if r:
        for i in r["images"]:
            word, info = process_images([i], message.from_user.id, message.from_user.first_name)

            seed_value = info.split(", Seed: ")[1].split(",")[0]
            caption = f"**[{message.from_user.first_name}](tg://user?id={message.from_user.id})**\n\n"
            for key, value in payload.items():
                caption += f"{key.capitalize()} - **{value}**\n"
            caption += f"Seed - **{seed_value}**\n"

            # Ensure caption is within the allowed length
            if len(caption) > 1024:
                caption = caption[:1021] + "..."

            message.reply_photo(photo=f"{IMAGE_PATH}/{word}.png", caption=caption)
            K.delete()
    else:
        message.reply_text("Failed to generate image. Please try again later.")
        K.delete()

@app.on_message(filters.command(["img"]))
def img2img(client, message):
    if not message.reply_to_message or not message.reply_to_message.photo:
        message.reply_text("reply to an image with \n`/img < prompt > ds:0-1.0`\n\nds stand for `Denoising_strength` parameter. Set that low (like 0.2) if you just want to slightly change things. defaults to 0.4")
        return

    msgs = message.text.split(" ", 1)
    print(msgs)

    if len(msgs) == 1:
        message.reply_text("""Format :\n/img < prompt >\nforce: < 0.1-1.0, default 0.3 >
                           """)
        return

    payload = parse_input(" ".join(msgs[1:]))
    print(payload)
    photo = message.reply_to_message.photo
    photo_file = app.download_media(photo)
    init_image = encode_file_to_base64(photo_file)
    os.remove(photo_file)  # Clean up downloaded image file

    payload["init_images"] = [init_image]

    K = message.reply_text("Please Wait 10-15 Seconds")
    r = call_api('sdapi/v1/img2img', payload)

    if r:
        for i in r["images"]:
            word, info = process_images([i], message.from_user.id, message.from_user.first_name)

            caption = f"**[{message.from_user.first_name}](tg://user?id={message.from_user.id})**\n\n"
            prompt = payload["prompt"]
            caption += f"**{prompt}**\n"

            message.reply_photo(photo=f"{IMAGE_PATH}/{word}.png", caption=caption)
            K.delete()
    else:
        message.reply_text("Failed to process image. Please try again later.")
        K.delete()

@app.on_message(filters.command(["getmodels"]))
async def get_models(client, message):
    try:
        response = requests.get(f"{SD_URL}/sdapi/v1/sd-models")
        response.raise_for_status()
        models_json = response.json()

        buttons = [
            [InlineKeyboardButton(model["title"], callback_data=model["model_name"])]
            for model in models_json
        ]
        await message.reply_text("Select a model [checkpoint] to use", reply_markup=InlineKeyboardMarkup(buttons))
    except requests.RequestException as e:
        await message.reply_text(f"Failed to get models: {e}")

@app.on_callback_query()
async def process_callback(client, callback_query):
    sd_model_checkpoint = callback_query.data
    options = {"sd_model_checkpoint": sd_model_checkpoint}

    try:
        response = requests.post(f"{SD_URL}/sdapi/v1/options", json=options)
        response.raise_for_status()
        await callback_query.message.reply_text(f"Checkpoint set to {sd_model_checkpoint}")
    except requests.RequestException as e:
        await callback_query.message.reply_text(f"Failed to set checkpoint: {e}")

# @app.on_message(filters.command(["start"], prefixes=["/", "!"]))
# async def start(client, message):
#     buttons = [[InlineKeyboardButton("Add to your group", url="https://t.me/gootmornbot?startgroup=true")]]
#     await message.reply_text("Hello!\nAsk me to imagine anything\n\n/draw text to image", reply_markup=InlineKeyboardMarkup(buttons))

user_interactions = {}

@app.on_message(filters.command(["user_stats"]))
def user_stats(client, message):
    stats = "User Interactions:\n\n"
    for user_id, info in user_interactions.items():
        stats += f"User: {info['username']} (ID: {user_id})\n"
        stats += f"Commands: {', '.join(info['commands'])}\n\n"
    
    message.reply_text(stats)

app.run()