├── .gitignore ├── README.md ├── assets ├── example1.jpg ├── example2.jpg └── example3.jpg ├── bot.py └── requirements.txt /.gitignore: -------------------------------------------------------------------------------- 1 | __pycache__/ 2 | .env -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # StableDiffusionTelegram 2 | StableDiffusionTelegram is a telegram bot that allows to generate images using the [Stable Diffusion](https://github.com/CompVis/stable-diffusion) AI from a telegram bot, in a much more comfortable and simple way. This bot can generate images from a text input or from an image with caption. In addition, given any response from the bot, a new try or a variation of the attempt can be generated. 3 | 4 | 5 | ## Try it 6 | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/11HvpaEJudoY5-KYU3GBWYjAMWoe2a21u) 7 | 8 | Try the bot with colab, create a token of [Hugging Face Hub](https://huggingface.co/docs/hub/security-tokens) and [Telegram](https://t.me/BotFather) and deploy the bot with colab. 9 | 10 | ## Installing 11 | 1. Install PyTorch (https://pytorch.org/get-started/). 12 | 13 | 2. Install requirements: 14 | ``` 15 | pip install -r requirements.txt 16 | ``` 17 | 18 | 3. Register in Hugging Face Hub and create a token (https://huggingface.co/docs/hub/security-tokens). 19 | 20 | 4. Use the token to login in huggingface-cli: 21 | ``` 22 | $ huggingface-cli login 23 | ``` 24 | 25 | 5. Talk to BotFather and create a bot (https://t.me/BotFather). 26 | 27 | 6. Create a .env file with the telegram token and the safe content option (if false, explicit content will be displayed, otherwise set to true). If you have memory problems, you can lower the width and height to lower numbers, such as 448 or 320. An example of .env is the following: 28 | ``` 29 | TG_TOKEN="YOUR_TOKEN_IS_HERE" 30 | SAFETY_CHECKER="false" 31 | HEIGHT="512" 32 | WIDTH="512" 33 | ``` 34 | 35 | 7. Run the bot 36 | ``` 37 | python bot.py 38 | ``` 39 | 40 | 41 | ## Examples 42 | Generating image from text | Generating a variation | Generating a new image from an user photo 43 | :-------------------------:|:-------------------------:|:-------------------------: 44 | ![](assets/example1.jpg) | ![](assets/example2.jpg) | ![](assets/example3.jpg) 45 | 46 | 47 | ## Credits 48 | 49 | * [Stable Diffusion](https://github.com/CompVis/stable-diffusion) 50 | * [Diffusers](https://github.com/huggingface/diffusers) 51 | * [python-telegram-bot](https://github.com/python-telegram-bot/python-telegram-bot) 52 | -------------------------------------------------------------------------------- /assets/example1.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/jossalgon/StableDiffusionTelegram/4def3953036af24f0c69dcbe92e90e8d1c7d9c5a/assets/example1.jpg -------------------------------------------------------------------------------- /assets/example2.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/jossalgon/StableDiffusionTelegram/4def3953036af24f0c69dcbe92e90e8d1c7d9c5a/assets/example2.jpg -------------------------------------------------------------------------------- /assets/example3.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/jossalgon/StableDiffusionTelegram/4def3953036af24f0c69dcbe92e90e8d1c7d9c5a/assets/example3.jpg -------------------------------------------------------------------------------- /bot.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import autocast 3 | from diffusers import StableDiffusionPipeline, StableDiffusionImg2ImgPipeline 4 | from PIL import Image 5 | 6 | import os 7 | from dotenv import load_dotenv 8 | from telegram import Update, InlineKeyboardButton, InlineKeyboardMarkup 9 | from telegram.ext import ApplicationBuilder, CallbackQueryHandler, ContextTypes, MessageHandler, filters 10 | from io import BytesIO 11 | import random 12 | 13 | load_dotenv() 14 | TG_TOKEN = os.getenv('TG_TOKEN') 15 | MODEL_DATA = os.getenv('MODEL_DATA', 'runwayml/stable-diffusion-v1-5') 16 | LOW_VRAM_MODE = (os.getenv('LOW_VRAM', 'true').lower() == 'true') 17 | USE_AUTH_TOKEN = (os.getenv('USE_AUTH_TOKEN', 'true').lower() == 'true') 18 | SAFETY_CHECKER = (os.getenv('SAFETY_CHECKER', 'true').lower() == 'true') 19 | HEIGHT = int(os.getenv('HEIGHT', '512')) 20 | WIDTH = int(os.getenv('WIDTH', '512')) 21 | NUM_INFERENCE_STEPS = int(os.getenv('NUM_INFERENCE_STEPS', '50')) 22 | STRENTH = float(os.getenv('STRENTH', '0.75')) 23 | GUIDANCE_SCALE = float(os.getenv('GUIDANCE_SCALE', '7.5')) 24 | 25 | revision = "fp16" if LOW_VRAM_MODE else None 26 | torch_dtype = torch.float16 if LOW_VRAM_MODE else None 27 | 28 | # load the text2img pipeline 29 | pipe = StableDiffusionPipeline.from_pretrained(MODEL_DATA, revision=revision, torch_dtype=torch_dtype, use_auth_token=USE_AUTH_TOKEN) 30 | pipe = pipe.to("cpu") 31 | 32 | # load the img2img pipeline 33 | img2imgPipe = StableDiffusionImg2ImgPipeline.from_pretrained(MODEL_DATA, revision=revision, torch_dtype=torch_dtype, use_auth_token=USE_AUTH_TOKEN) 34 | img2imgPipe = img2imgPipe.to("cpu") 35 | 36 | # disable safety checker if wanted 37 | def dummy_checker(images, **kwargs): return images, False 38 | if not SAFETY_CHECKER: 39 | pipe.safety_checker = dummy_checker 40 | img2imgPipe.safety_checker = dummy_checker 41 | 42 | 43 | def image_to_bytes(image): 44 | bio = BytesIO() 45 | bio.name = 'image.jpeg' 46 | image.save(bio, 'JPEG') 47 | bio.seek(0) 48 | return bio 49 | 50 | def get_try_again_markup(): 51 | keyboard = [[InlineKeyboardButton("Try again", callback_data="TRYAGAIN"), InlineKeyboardButton("Variations", callback_data="VARIATIONS")]] 52 | reply_markup = InlineKeyboardMarkup(keyboard) 53 | return reply_markup 54 | 55 | 56 | def generate_image(prompt, seed=None, height=HEIGHT, width=WIDTH, num_inference_steps=NUM_INFERENCE_STEPS, strength=STRENTH, guidance_scale=GUIDANCE_SCALE, photo=None): 57 | seed = seed if seed is not None else random.randint(1, 10000) 58 | generator = torch.cuda.manual_seed_all(seed) 59 | 60 | if photo is not None: 61 | pipe.to("cpu") 62 | img2imgPipe.to("cuda") 63 | init_image = Image.open(BytesIO(photo)).convert("RGB") 64 | init_image = init_image.resize((height, width)) 65 | with autocast("cuda"): 66 | image = img2imgPipe(prompt=[prompt], init_image=init_image, 67 | generator=generator, 68 | strength=strength, 69 | guidance_scale=guidance_scale, 70 | num_inference_steps=num_inference_steps)["images"][0] 71 | else: 72 | pipe.to("cuda") 73 | img2imgPipe.to("cpu") 74 | with autocast("cuda"): 75 | image = pipe(prompt=[prompt], 76 | generator=generator, 77 | strength=strength, 78 | height=height, 79 | width=width, 80 | guidance_scale=guidance_scale, 81 | num_inference_steps=num_inference_steps)["images"][0] 82 | return image, seed 83 | 84 | 85 | async def generate_and_send_photo(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: 86 | progress_msg = await update.message.reply_text("Generating image...", reply_to_message_id=update.message.message_id) 87 | im, seed = generate_image(prompt=update.message.text) 88 | await context.bot.delete_message(chat_id=progress_msg.chat_id, message_id=progress_msg.message_id) 89 | await context.bot.send_photo(update.effective_user.id, image_to_bytes(im), caption=f'"{update.message.text}" (Seed: {seed})', reply_markup=get_try_again_markup(), reply_to_message_id=update.message.message_id) 90 | 91 | async def generate_and_send_photo_from_photo(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: 92 | if update.message.caption is None: 93 | await update.message.reply_text("The photo must contain a text in the caption", reply_to_message_id=update.message.message_id) 94 | return 95 | progress_msg = await update.message.reply_text("Generating image...", reply_to_message_id=update.message.message_id) 96 | photo_file = await update.message.photo[-1].get_file() 97 | photo = await photo_file.download_as_bytearray() 98 | im, seed = generate_image(prompt=update.message.caption, photo=photo) 99 | await context.bot.delete_message(chat_id=progress_msg.chat_id, message_id=progress_msg.message_id) 100 | await context.bot.send_photo(update.effective_user.id, image_to_bytes(im), caption=f'"{update.message.caption}" (Seed: {seed})', reply_markup=get_try_again_markup(), reply_to_message_id=update.message.message_id) 101 | 102 | 103 | async def button(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: 104 | query = update.callback_query 105 | replied_message = query.message.reply_to_message 106 | 107 | await query.answer() 108 | progress_msg = await query.message.reply_text("Generating image...", reply_to_message_id=replied_message.message_id) 109 | 110 | if query.data == "TRYAGAIN": 111 | if replied_message.photo is not None and len(replied_message.photo) > 0 and replied_message.caption is not None: 112 | photo_file = await replied_message.photo[-1].get_file() 113 | photo = await photo_file.download_as_bytearray() 114 | prompt = replied_message.caption 115 | im, seed = generate_image(prompt, photo=photo) 116 | else: 117 | prompt = replied_message.text 118 | im, seed = generate_image(prompt) 119 | elif query.data == "VARIATIONS": 120 | photo_file = await query.message.photo[-1].get_file() 121 | photo = await photo_file.download_as_bytearray() 122 | prompt = replied_message.text if replied_message.text is not None else replied_message.caption 123 | im, seed = generate_image(prompt, photo=photo) 124 | await context.bot.delete_message(chat_id=progress_msg.chat_id, message_id=progress_msg.message_id) 125 | await context.bot.send_photo(update.effective_user.id, image_to_bytes(im), caption=f'"{prompt}" (Seed: {seed})', reply_markup=get_try_again_markup(), reply_to_message_id=replied_message.message_id) 126 | 127 | 128 | 129 | app = ApplicationBuilder().token(TG_TOKEN).build() 130 | 131 | app.add_handler(MessageHandler(filters.TEXT & ~filters.COMMAND, generate_and_send_photo)) 132 | app.add_handler(MessageHandler(filters.PHOTO, generate_and_send_photo_from_photo)) 133 | app.add_handler(CallbackQueryHandler(button)) 134 | 135 | app.run_polling() 136 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | diffusers==0.7.2 2 | transformers==4.24.0 3 | scipy==1.7.3 4 | ftfy==6.1.1 5 | python-telegram-bot>=20.0a2 6 | python-dotenv>=0.20.0 7 | huggingface_hub==0.10.1 8 | --------------------------------------------------------------------------------