mirror of
				https://github.com/5shekel/stable-diffusion-telegram-bot.git
				synced 2024-05-22 19:33:14 +03:00 
			
		
		
		
	add steps and nicer filename
This commit is contained in:
		
							parent
							
								
									9bce6f989e
								
							
						
					
					
						commit
						095e418da9
					
				
							
								
								
									
										90
									
								
								main.py
									
									
									
									
									
								
							
							
						
						
									
										90
									
								
								main.py
									
									
									
									
									
								
							| @ -2,7 +2,7 @@ import json | ||||
| import requests | ||||
| import io | ||||
| import os | ||||
| import random | ||||
| import uuid | ||||
| import base64 | ||||
| from PIL import Image, PngImagePlugin | ||||
| from pyrogram import Client, filters | ||||
| @ -26,55 +26,81 @@ SD_URL = os.environ.get("SD_URL", None) | ||||
| print(SD_URL) | ||||
| app = Client("stable", api_id=API_ID, api_hash=API_HASH, bot_token=TOKEN) | ||||
| 
 | ||||
| #default params | ||||
| steps_value_default = 40 | ||||
| 
 | ||||
| def slice_positive_negative(string): | ||||
|     delimiter = "ng:" | ||||
|     if delimiter in string: | ||||
|         index = string.index(delimiter) | ||||
|         positive = string[:index].rstrip() | ||||
|         negative = string[index + len(delimiter) :].lstrip() | ||||
|         return positive, negative | ||||
| def process_input_string(string): | ||||
|     ng_delimiter = "ng:" | ||||
|     steps_delimiter = "steps:" | ||||
|      | ||||
|     ng_index = string.find(ng_delimiter) | ||||
|     steps_index = string.find(steps_delimiter) | ||||
|      | ||||
|     if ng_index != -1 and steps_index != -1: | ||||
|         if ng_index < steps_index: | ||||
|             positive = string[:ng_index].strip() | ||||
|             negative = string[ng_index + len(ng_delimiter):steps_index].strip() | ||||
|             steps_str = string[steps_index + len(steps_delimiter):].strip().split()[0] | ||||
|         else: | ||||
|         return string, "null"  # if we are missing the ng bad thing happen.... | ||||
|         #POSSIBLE NG [[[[[[[[[[lowres, low resolution, bad resolution, bad, poor quality, bad quality, blurry, bad art, incomplete, logo, signature, text, jpeg artifacts, compression artifacts,child, childish]]]]]]]]]] | ||||
|             positive = string[:steps_index].strip() | ||||
|             negative = string[steps_index + len(steps_delimiter):ng_index].strip() | ||||
|             steps_str = string[ng_index + len(ng_delimiter):].strip().split()[0] | ||||
|     elif ng_index != -1: | ||||
|         positive = string[:ng_index].strip() | ||||
|         negative = string[ng_index + len(ng_delimiter):].strip() | ||||
|         steps_str = None | ||||
|     elif steps_index != -1: | ||||
|         positive = string[:steps_index].strip() | ||||
|         negative = None | ||||
|         steps_str = string[steps_index + len(steps_delimiter):].strip().split()[0] | ||||
|     else: | ||||
|         positive = string.strip() | ||||
|         negative = None | ||||
|         steps_str = None | ||||
| 
 | ||||
|     try: | ||||
|         steps_value = int(steps_str) | ||||
|         #limit steps to range | ||||
|         if not 1 <= steps_value <= 70: | ||||
|             steps_value = steps_value_default | ||||
|     except (ValueError, TypeError): | ||||
|         steps_value = None | ||||
|      | ||||
| 
 | ||||
|     return positive, negative, steps_value | ||||
| 
 | ||||
| @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) >" | ||||
|             "Format :\n/draw < text to image >\nng: < negative (optional) >\nsteps: < steps value (1-70, optional) >" | ||||
|         ) | ||||
|         return | ||||
| 
 | ||||
|     msg = slice_positive_negative(msgs[1]) | ||||
|     print(msg) | ||||
|     positive, negative, steps_value = process_input_string(msgs[1]) | ||||
|     payload = { | ||||
|         "prompt": msg[0], | ||||
|         "negative_prompt": msg[1], | ||||
|         "prompt": positive, | ||||
|     } | ||||
|     if negative is not None: | ||||
|         payload["negative_prompt"] = negative | ||||
|     if steps_value is not None: | ||||
|         payload["steps"] = steps_value | ||||
| 
 | ||||
|     print(payload) | ||||
| 
 | ||||
|     # The rest of the draw function remains unchanged | ||||
| 
 | ||||
| 
 | ||||
|     K = message.reply_text("Please Wait 10-15 Second") | ||||
|     r = requests.post(url=f"{SD_URL}/sdapi/v1/txt2img", json=payload).json() | ||||
| 
 | ||||
|     def genr(): | ||||
|         chars = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ" | ||||
|         chars1 = "1234564890" | ||||
|         gen1 = random.choice(chars) | ||||
|         gen2 = random.choice(chars) | ||||
|         gen3 = random.choice(chars1) | ||||
|         gen4 = random.choice(chars) | ||||
|         gen5 = random.choice(chars) | ||||
|         gen6 = random.choice(chars) | ||||
|         gen7 = random.choice(chars1) | ||||
|         gen8 = random.choice(chars) | ||||
|         gen9 = random.choice(chars) | ||||
|         gen10 = random.choice(chars1) | ||||
|         return f"{message.from_user.id}-MOE{gen1}{gen2}{gen3}{gen4}{gen5}{gen6}{gen7}{gen8}{gen9}{gen10}" | ||||
|         unique_id = str(uuid.uuid4())[:7] | ||||
|         return f"{message.from_user.first_name}-{unique_id}" | ||||
| 
 | ||||
|     word = genr() | ||||
| 
 | ||||
|     for i in r["images"]: | ||||
|         image = Image.open(io.BytesIO(base64.b64decode(i.split(",", 1)[0]))) | ||||
| 
 | ||||
| @ -87,8 +113,14 @@ def draw(client, message): | ||||
| 
 | ||||
|         message.reply_photo( | ||||
|         photo=f"{word}.png", | ||||
|             caption=f"Prompt - **{msg[0]}**\nNegative Prompt - **{msg[1]}**\n**[{message.from_user.first_name}-Kun](tg://user?id={message.from_user.id})**", | ||||
|         caption=( | ||||
|             f"Prompt - **{positive}**\n" | ||||
|             f"Negative Prompt - **{negative if negative is not None else 'None'}**\n" | ||||
|             f"Steps - **{steps_value if steps_value != steps_value_default else 'Default'}**\n" | ||||
|             f"**[{message.from_user.first_name}-Kun](tg://user?id={message.from_user.id})**" | ||||
|         ), | ||||
| ) | ||||
| 
 | ||||
|         # os.remove(f"{word}.png") | ||||
|         K.delete() | ||||
| 
 | ||||
|  | ||||
		Loading…
	
	
			
			x
			
			
		
	
		Reference in New Issue
	
	Block a user