size limit

main.py

@@ -17,6 +17,10 @@ API_HASH = os.environ.get("API_HASH")
 TOKEN = os.environ.get("TOKEN_givemtxt2img")
 SD_URL = os.environ.get("SD_URL")
 
+# Ensure all required environment variables are loaded
+if not all([API_ID, API_HASH, TOKEN, SD_URL]):
+    raise EnvironmentError("Missing one or more required environment variables: API_ID, API_HASH, TOKEN, SD_URL")
+
 app = Client("stable", api_id=API_ID, api_hash=API_HASH, bot_token=TOKEN)
 IMAGE_PATH = 'images'
 
@@ -25,6 +29,7 @@ os.makedirs(IMAGE_PATH, exist_ok=True)
 
 # Model-specific embeddings for negative prompts
 model_negative_prompts = {
+    "coloringPage_v10": "fake",
     "Anything-Diffusion": "",
     "Deliberate": "",
     "Dreamshaper": "",
@@ -39,10 +44,12 @@ 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))
 
+
 # Set default payload values
 default_payload = {
     "prompt": "",
@@ -62,12 +69,34 @@ default_payload = {
     "override_settings_restore_afterwards": True,
 }
 
+
 def update_negative_prompt(model_name):
+    """Update the negative prompt for a given model."""
     if model_name in model_negative_prompts:
         suffix = model_negative_prompts[model_name]
         default_payload["negative_prompt"] += f", {suffix}"
 
+
+def update_resolution(model_name):
+    """Update resolution based on the selected model."""
+    if model_name == "Juggernaut-XL_v9_RunDiffusionPhoto_v2":
+        default_payload["width"] = 832
+        default_payload["height"] = 1216
+    else:
+        default_payload["width"] = 512
+        default_payload["height"] = 512
+
+
+def update_cfg_scale(model_name):
+    """Update CFG scale based on the selected model."""
+    if model_name == "Juggernaut-XL_v9_RunDiffusionPhoto_v2":
+        default_payload["cfg_scale"] = 1
+    else:
+        default_payload["cfg_scale"] = 7
+
+
 def parse_input(input_string):
+    """Parse the input string and create a payload."""
     payload = default_payload.copy()
     prompt = []
     include_info = "info:" in input_string
@@ -129,11 +158,6 @@ def parse_input(input_string):
         else:
             prompt.append(f"{key}:{value}")
 
-        # Adjust dimensions for specific model
-        if "Juggernaut-XL_v9_RunDiffusionPhoto_v2" in input_string:
-            payload["width"] = 823
-            payload["height"] = 1216
-            
         last_index = value_end_index
 
     payload["prompt"] = " ".join(prompt).strip()
@@ -147,18 +171,13 @@ def parse_input(input_string):
     return payload, include_info
 
 
-
 def create_caption(payload, user_name, user_id, info, include_info):
+    """Create a caption for the generated image."""
     caption = f"**[{user_name}](tg://user?id={user_id})**\n\n"
     prompt = payload["prompt"]
 
-    # Define a regular expression pattern to match the seed value
     seed_pattern = r"Seed: (\d+)"
-
-    # Search for the pattern in the info string
     match = re.search(seed_pattern, info)
-
-    # Check if a match was found and extract the seed value
     if match:
         seed_value = match.group(1)
         caption += f"**{seed_value}**\n"
@@ -177,15 +196,18 @@ def create_caption(payload, user_name, user_id, info, include_info):
 
 
 def call_api(api_endpoint, payload):
+    """Call the API with the provided 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
+        return {"error": str(e)}
+
 
 def process_images(images, user_id, user_name):
+    """Process and save generated images."""
     def generate_unique_name():
         unique_id = str(uuid.uuid4())[:7]
         return f"{user_name}-{unique_id}"
@@ -202,11 +224,16 @@ def process_images(images, user_id, user_name):
         pnginfo.add_text("parameters", response2.json().get("info"))
         image.save(f"{IMAGE_PATH}/{word}.png", pnginfo=pnginfo)
 
+        # Save as JPG
+        jpg_path = f"{IMAGE_PATH}/{word}.jpg"
+        image.convert("RGB").save(jpg_path, "JPEG")
+        
         return word, response2.json().get("info")
 
 
 @app.on_message(filters.command(["draw"]))
 def draw(client, message):
+    """Handle /draw command to generate images from text prompts."""
     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) >")
@@ -214,7 +241,6 @@ def draw(client, message):
 
     payload, include_info = parse_input(msgs[1])
 
-    # Check if xds is used in the payload
     if "xds" in msgs[1].lower():
         message.reply_text("`xds` key cannot be used in the `/draw` command. Use `/img` instead.")
         return
@@ -222,25 +248,28 @@ def draw(client, message):
     K = message.reply_text("Please Wait 10-15 Seconds")
     r = call_api('sdapi/v1/txt2img', payload)
 
-    if r:
+    if r and "images" in r:
         for i in r["images"]:
             word, info = process_images([i], message.from_user.id, message.from_user.first_name)
             caption = create_caption(payload, message.from_user.first_name, message.from_user.id, info, include_info)
-            message.reply_photo(photo=f"{IMAGE_PATH}/{word}.png", caption=caption)
+            message.reply_photo(photo=f"{IMAGE_PATH}/{word}.jpg", caption=caption)
         K.delete()
     else:
-        message.reply_text("Failed to generate image. Please try again later.")
+        error_message = r.get("error", "Failed to generate image. Please try again later.")
+        message.reply_text(error_message)
         K.delete()
 
+
 @app.on_message(filters.command(["img"]))
 def img2img(client, message):
+    """Handle /img command to generate images from existing images."""
     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 stands for `Denoising_strength` parameter. Set that low (like 0.2) if you just want to slightly change things. defaults to 0.35\n\nExample: `/img murder on the dance floor ds:0.2`")
         return
 
     msgs = message.text.split(" ", 1)
     if len(msgs) == 1:
-        message.reply_text("dont FAIL in life")
+        message.reply_text("Don't FAIL in life")
         return
 
     payload, include_info = parse_input(msgs[1])
@@ -254,19 +283,21 @@ def img2img(client, message):
     K = message.reply_text("Please Wait 10-15 Seconds")
     r = call_api('sdapi/v1/img2img', payload)
 
-    if r:
+    if r and "images" in r:
         for i in r["images"]:
             word, info = process_images([i], message.from_user.id, message.from_user.first_name)
             caption = create_caption(payload, message.from_user.first_name, message.from_user.id, info, include_info)
-            message.reply_photo(photo=f"{IMAGE_PATH}/{word}.png", caption=caption)
+            message.reply_photo(photo=f"{IMAGE_PATH}/{word}.jpg", caption=caption)
         K.delete()
     else:
-        message.reply_text("Failed to process image. Please try again later.")
+        error_message = r.get("error", "Failed to process image. Please try again later.")
+        message.reply_text(error_message)
         K.delete()
 
 
 @app.on_message(filters.command(["getmodels"]))
 async def get_models(client, message):
+    """Handle /getmodels command to list available models."""
     try:
         response = requests.get(f"{SD_URL}/sdapi/v1/sd-models")
         response.raise_for_status()
@@ -279,8 +310,10 @@ async def get_models(client, message):
     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):
+    """Process model selection from callback queries."""
     sd_model_checkpoint = callback_query.data
     options = {"sd_model_checkpoint": sd_model_checkpoint}
 
@@ -288,16 +321,19 @@ async def process_callback(client, callback_query):
         response = requests.post(f"{SD_URL}/sdapi/v1/options", json=options)
         response.raise_for_status()
 
-        # Update the negative prompt based on the selected model
         update_negative_prompt(sd_model_checkpoint)
+        update_resolution(sd_model_checkpoint)
+        update_cfg_scale(sd_model_checkpoint)
         
         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}")
         print(f"Error setting checkpoint: {e}")
 
+
 @app.on_message(filters.command(["info_sd_bot"]))
 async def info(client, message):
+    """Provide information about the bot's commands and options."""
     await message.reply_text("""
 **Stable Diffusion Bot Commands and Options:**