refactor with rate limits and troll check

app.py

@@ -147,15 +147,15 @@ with gr.Blocks() as demo:
         with gr.Column():
             gr.Markdown("""
             <div style="text-align: center;">
-                Generate an image using a color palette.  If you choose to include an image (optional) the subject and style will be used as a reference. 
+                Generate an image using a color palette.  If you must include an image and text prompt, the subject and style will be used as a reference. 
                 The colors of the image will also be incorporated, along with the colors from the colors list. A color list is always required but one has been provided.
                 </div>
             """)
             reference_image = gr.Image(type='pil', label="Reference Image")     
             colors = gr.Textbox(label="Colors", placeholder="Enter up to 10 colors as hex values, e.g., #00FF00,#FCF2AB", max_lines=1)
-            with gr.Accordion("Optional Prompt", open=False):
-                prompt = gr.Textbox(label="Text", placeholder="Enter a text prompt (1-1024 characters)", max_lines=4)
-                gr.Button("Generate Prompt").click(generate_nova_prompt, outputs=prompt)
+            
+            prompt = gr.Textbox(label="Text", placeholder="Enter a text prompt (1-1024 characters)", max_lines=4)
+            gr.Button("Generate Prompt").click(generate_nova_prompt, outputs=prompt)
             error_box = gr.Markdown(visible=False, label="Error", elem_classes="center-markdown")
             output = gr.Image()
             with gr.Accordion("Advanced Options", open=False):

functions.py

@@ -5,6 +5,7 @@ import gradio as gr
 from PIL import Image
 from generate import *
 from typing import Dict, Any
+from processImage import process_and_encode_image
 
 def display_image(image_bytes):
     if isinstance(image_bytes, str):
@@ -56,6 +57,12 @@ def build_request(task_type, params, height=1024, width=1024, quality="standard"
         )
     })
 
+def check_return(result):
+    if not isinstance(result, bytes):
+        return None, gr.update(visible=True, value=result)
+        
+    return Image.open(io.BytesIO(result)), gr.update(visible=False)
+
 
 def text_to_image(prompt, negative_text=None, height=1024, width=1024, quality="standard", cfg_scale=8.0, seed=0):
     text_to_image_params = {"text": prompt,
@@ -63,15 +70,16 @@ def text_to_image(prompt, negative_text=None, height=1024, width=1024, quality="
                             }
     
     body = build_request("TEXT_IMAGE", text_to_image_params, height, width, quality, cfg_scale, seed)
-    image_bytes = generate_image(body)
-    return display_image(image_bytes)
+    result = generate_image(body)
+    return check_return(result)
+    
 
 def inpainting(image, mask_prompt=None, mask_image=None, text=None, negative_text=None, height=1024, width=1024, quality="standard", cfg_scale=8.0, seed=0):
     images = process_images(primary=image, secondary=None)
     
     for value in images.values():
-        if isinstance(value, str) and "Not Appropriate" in value:
-            return None, gr.update(visible=True, value="Image <b>Not Appropriate</b>")
+        if len(value) < 200:
+            return None, gr.update(visible=True, value=value)
     # Prepare the inPaintingParams dictionary
     if mask_prompt and mask_image:
         raise ValueError("You must specify either maskPrompt or maskImage, but not both.")
@@ -87,13 +95,15 @@ def inpainting(image, mask_prompt=None, mask_image=None, text=None, negative_tex
     }
 
     body = build_request("INPAINTING", in_painting_params, height, width, quality, cfg_scale, seed)
-    return display_image(generate_image(body))
+    result = generate_image(body)
+    
+    return check_return(result)
 
 def outpainting(image, mask_prompt=None, mask_image=None, text=None, negative_text=None, outpainting_mode="DEFAULT", height=1024, width=1024, quality="standard", cfg_scale=8.0, seed=0):
     images = process_images(primary=image, secondary=None)
     for value in images.values():
-        if isinstance(value, str) and "Not Appropriate" in value:
-            return None, gr.update(visible=True, value="Image <b>Not Appropriate</b>")
+        if len(value) < 200:
+            return None, gr.update(visible=True, value=value)
 
     if mask_prompt and mask_image:
         raise ValueError("You must specify either maskPrompt or maskImage, but not both.")
@@ -111,13 +121,19 @@ def outpainting(image, mask_prompt=None, mask_image=None, text=None, negative_te
     }
 
     body = build_request("OUTPAINTING", out_painting_params, height, width, quality, cfg_scale, seed)
-    return display_image(generate_image(body))
+    result = generate_image(body)
+    
+    return check_return(result)
 
 def image_variation(images, text=None, negative_text=None, similarity_strength=0.5, height=1024, width=1024, quality="standard", cfg_scale=8.0, seed=0):
     encoded_images = []
     for image_path in images:
         with open(image_path, "rb") as image_file:
-            encoded_images.append(process_and_encode_image(image_file))
+            value = process_and_encode_image(image_file)
+            
+            if len(value) < 200:
+                return None, gr.update(visible=True, value=value)
+            encoded_images.append(value)
 
     # Prepare the imageVariationParams dictionary
     image_variation_params = {
@@ -127,30 +143,34 @@ def image_variation(images, text=None, negative_text=None, similarity_strength=0
     }
 
     body = build_request("IMAGE_VARIATION", image_variation_params, height, width, quality, cfg_scale, seed)
-    return display_image(generate_image(body))
+    result = generate_image(body)
+    
+    return check_return(result)
 
 def image_conditioning(condition_image, text, negative_text=None, control_mode="CANNY_EDGE", control_strength=0.7, height=1024, width=1024, quality="standard", cfg_scale=8.0, seed=0):
     condition_image_encoded = process_images(primary=condition_image)
     for value in condition_image_encoded.values():
-        if isinstance(value, str) and "Not Appropriate" in value:
-            return None, gr.update(visible=True, value="Image <b>Not Appropriate</b>")
+        if len(value) < 200:
+            return None, gr.update(visible=True, value=value)
     # Prepare the textToImageParams dictionary
     text_to_image_params = {
         "text": text,
         "controlMode": control_mode,
         "controlStrength": control_strength,
-        **condition_image_encoded,
+        "conditionImage": condition_image_encoded.get('image'),
         **({"negativeText": negative_text} if negative_text not in [None, ""] else {})
     }
     body = build_request("TEXT_IMAGE", text_to_image_params, height, width, quality, cfg_scale, seed)
-    return display_image(generate_image(body))
+    result = generate_image(body)
+    
+    return check_return(result)
 
 def color_guided_content(text=None, reference_image=None, negative_text=None, colors=None, height=1024, width=1024, quality="standard", cfg_scale=8.0, seed=0):
     # Encode the reference image if provided
     reference_image_encoded = process_images(primary=reference_image)
     for value in reference_image_encoded.values():
-        if isinstance(value, str) and "Not Appropriate" in value:
-            return None, gr.update(visible=True, value="Image <b>Not Appropriate</b>")
+        if len(value) < 200:
+            return None, gr.update(visible=True, value=value)
         
     if not colors:
         colors = "#FF5733,#33FF57,#3357FF,#FF33A1,#33FFF5,#FF8C33,#8C33FF,#33FF8C,#FF3333,#33A1FF"
@@ -158,24 +178,30 @@ def color_guided_content(text=None, reference_image=None, negative_text=None, co
     color_guided_generation_params = {
         "text": text,
         "colors": colors.split(','),
-        **reference_image_encoded,
+        "referenceImage": reference_image_encoded.get('image'),
         **({"negativeText": negative_text} if negative_text not in [None, ""] else {})
     }
 
     body = build_request("COLOR_GUIDED_GENERATION", color_guided_generation_params, height, width, quality, cfg_scale, seed)
-    return display_image(generate_image(body))
+    result = generate_image(body)
+    
+    return check_return(result)
 
 def background_removal(image):
-    input_image = process_and_encode_image(image)
+    input_image = process_images(primary=image)
     for value in input_image.values():
-        if isinstance(value, str) and "Not Appropriate" in value:
-            return None, gr.update(visible=True, value="Image <b>Not Appropriate</b>")
+        if len(value) < 200:
+            return None, gr.update(visible=True, value=value)
         
     body = json.dumps({
         "taskType": "BACKGROUND_REMOVAL",
-        "backgroundRemovalParams": {"image": input_image}
+        "backgroundRemovalParams": {
+            "image": input_image.get('image')  
+        }
     })
-    return display_image(generate_image(body))
+    result = generate_image(body)
+    
+    return check_return(result)
 
 def generate_nova_prompt():
     

generate.py

@@ -3,14 +3,9 @@ import base64
 import boto3
 import json
 import logging
-import io
-import time
-import requests
 from datetime import datetime
 from dotenv import load_dotenv
-from PIL import Image
 from functools import wraps
-from dataclasses import dataclass
 from botocore.config import Config
 from botocore.exceptions import ClientError
 
@@ -35,121 +30,13 @@ def handle_bedrock_errors(func):
             raise ImageError(f"Unexpected error: {str(err)}")
     return wrapper
 
-@dataclass
-class ImageConfig:
-    min_size: int = 320
-    max_size: int = 4096
-    max_pixels: int = 4194304
-    quality: str = "standard"
-    format: str = "PNG"
-
-config = ImageConfig()
-
-model_id = 'amazon.nova-canvas-v1:0'
 aws_id = os.getenv('AWS_ID')
 aws_secret = os.getenv('AWS_SECRET')
-token = os.environ.get("HF_TOKEN")
-headers = {"Authorization": f"Bearer {token}", "x-use-cache": "0", 'Content-Type': 'application/json'}
 nova_image_bucket='nova-image-data'
 bucket_region='us-west-2'
-
-class ImageProcessor:
-    def __init__(self, image):
-        self.image = self._open_image(image)
-        
-    def _open_image(self, image):
-        """Convert input to PIL Image if necessary."""
-        if image is None:
-            raise ValueError("Input image is required.")
-        return Image.open(image) if not isinstance(image, Image.Image) else image
-    
-    def _check_nsfw(self, attempts=1):
-        """Check if image is NSFW using Hugging Face API."""
-        API_URL = "https://api-inference.huggingface.co/models/Falconsai/nsfw_image_detection"
-        
-        # Prepare image data
-        temp_buffer = io.BytesIO()
-        self.image.save(temp_buffer, format='PNG')
-        temp_buffer.seek(0)
-        
-        try:
-            response = requests.request("POST", API_URL, headers=headers, data=temp_buffer.getvalue())
-            json_response = json.loads(response.content.decode("utf-8"))
-            print(json_response)
-            if "error" in json_response:
-                if attempts > 30:
-                    raise ImageError("NSFW check failed after multiple attempts")
-                time.sleep(json_response["estimated_time"])
-                return self._check_nsfw(attempts + 1)
-            
-            nsfw_score = next((item['score'] for item in json_response if item['label'] == 'nsfw'), 0)
-            print(f"NSFW Score: {nsfw_score}")
-            
-            if nsfw_score > 0.1:
-                return None
-                  
-            return self
-            
-        except json.JSONDecodeError as e:
-            raise ImageError(f"NSFW check failed: Invalid response format - {str(e)}")
-        except Exception as e:
-            if attempts > 30:
-                raise ImageError("NSFW check failed after multiple attempts")
-            return self._check_nsfw(attempts + 1)
-    
-    def _convert_color_mode(self):
-        """Handle color mode conversion."""
-        if self.image.mode not in ('RGB', 'RGBA'):
-            self.image = self.image.convert('RGB')
-        elif self.image.mode == 'RGBA':
-            background = Image.new('RGB', self.image.size, (255, 255, 255))
-            background.paste(self.image, mask=self.image.split()[3])
-            self.image = background
-        return self
-    
-    def _resize_for_pixels(self, max_pixels):
-        """Resize image to meet pixel limit."""
-        current_pixels = self.image.width * self.image.height
-        if current_pixels > max_pixels:
-            aspect_ratio = self.image.width / self.image.height
-            if aspect_ratio > 1:
-                new_width = int((max_pixels * aspect_ratio) ** 0.5)
-                new_height = int(new_width / aspect_ratio)
-            else:
-                new_height = int((max_pixels / aspect_ratio) ** 0.5)
-                new_width = int(new_height * aspect_ratio)
-            self.image = self.image.resize((new_width, new_height), Image.LANCZOS)
-        return self
-    
-    def _ensure_dimensions(self, min_size=320, max_size=4096):
-        if (self.image.width < min_size or 
-            self.image.width > max_size or 
-            self.image.height < min_size or 
-            self.image.height > max_size):
-            
-            new_width = min(max(self.image.width, min_size), max_size)
-            new_height = min(max(self.image.height, min_size), max_size)
-            self.image = self.image.resize((new_width, new_height), Image.LANCZOS)
-        
-        return self
-    
-    def encode(self):
-        image_bytes = io.BytesIO()
-        self.image.save(image_bytes, format='PNG', optimize=True)
-        return base64.b64encode(image_bytes.getvalue()).decode('utf8')
-        
-    def process(self, min_size=320, max_size=4096, max_pixels=4194304):
-        """Process image with all necessary transformations."""
-        result = (self
-            ._convert_color_mode()
-            ._resize_for_pixels(max_pixels)
-            ._ensure_dimensions(min_size, max_size)
-            ._check_nsfw())  # Add NSFW check before encoding
-    
-        if result is None:
-            raise ImageError("Image <b>Not Appropriate</b>")
-            
-        return result.encode()
+rate_limit_message = """<div style='text-align: center;'>{} rate limit exceeded. Check back later, use the 
+            <a href='https://docs.aws.amazon.com/bedrock/latest/userguide/playgrounds.html'>Bedrock Playground</a> or
+            try it out without an AWS account on <a href='https://partyrock.aws/'>PartyRock</a>.</div>"""
 
 # Function to generate an image using Amazon Nova Canvas model
 class BedrockClient:
@@ -282,16 +169,12 @@ def check_rate_limit(body):
     
     # Check limits based on quality
     if quality == 'premium':
-        if len(rate_data['premium']) >= 4:   
-            raise ImageError("""<div style='text-align: center;'>Premium rate limit exceeded. Check back later, use the 
-            <a href='https://docs.aws.amazon.com/bedrock/latest/userguide/playgrounds.html'>Bedrock Playground</a> or
-            try it out without an AWS account on <a href='https://partyrock.aws/'>PartyRock</a>.</div>""")
+        if len(rate_data['premium']) >= 3:   
+            raise ImageError(rate_limit_message.format('Premium'))
         rate_data['premium'].append(current_time)
     else:  # standard
-        if len(rate_data['standard']) >= 8:
-            raise ImageError("""<div style='text-align: center;'>Standard rate limit exceeded. Check back later, use the 
-            <a href='https://docs.aws.amazon.com/bedrock/latest/userguide/playgrounds.html'>Bedrock Playground</a> or
-            try it out without an AWS account on <a href='https://partyrock.aws/'>PartyRock</a>.</div>""")
+        if len(rate_data['standard']) >= 6:
+            raise ImageError(rate_limit_message.format('Standard'))
         rate_data['standard'].append(current_time)
     
     # Update rate limit file
@@ -303,14 +186,6 @@ def check_rate_limit(body):
     )
     
 
-def process_and_encode_image(image, **kwargs):
-    """Process and encode image with default parameters."""
-    try:
-        image = ImageProcessor(image).process(**kwargs)
-        return image
-    except ImageError as e:
-        return str(e)
-
 def generate_image(body):
     """Generate image using Bedrock service."""
     try:

processImage.py

@@ -0,0 +1,134 @@
+import os
+import base64
+import json
+import io
+import time
+import requests
+from dotenv import load_dotenv
+from PIL import Image
+from dataclasses import dataclass
+
+load_dotenv()
+# Move custom exceptions to the top
+class ImageError(Exception):
+    def __init__(self, message):
+        self.message = message
+
+@dataclass
+class ImageConfig:
+    min_size: int = 320
+    max_size: int = 4096
+    max_pixels: int = 4194304
+    quality: str = "standard"
+    format: str = "PNG"
+
+config = ImageConfig()
+
+token = os.environ.get("HF_TOKEN")
+headers = {"Authorization": f"Bearer {token}", "x-use-cache": "0", 'Content-Type': 'application/json'}
+
+class ImageProcessor:
+    def __init__(self, image):
+        self.image = self._open_image(image)
+        
+    def _open_image(self, image):
+        """Convert input to PIL Image if necessary."""
+        if image is None:
+            raise ValueError("Input image is required.")
+        return Image.open(image) if not isinstance(image, Image.Image) else image
+    
+    def _check_nsfw(self, attempts=1):
+        """Check if image is NSFW using Hugging Face API."""
+        API_URL = "https://api-inference.huggingface.co/models/Falconsai/nsfw_image_detection"
+        
+        # Prepare image data
+        temp_buffer = io.BytesIO()
+        self.image.save(temp_buffer, format='PNG')
+        temp_buffer.seek(0)
+        
+        try:
+            response = requests.request("POST", API_URL, headers=headers, data=temp_buffer.getvalue())
+            json_response = json.loads(response.content.decode("utf-8"))
+            print(json_response)
+            if "error" in json_response:
+                if attempts > 30:
+                    raise ImageError("NSFW check failed after multiple attempts")
+                time.sleep(json_response["estimated_time"])
+                return self._check_nsfw(attempts + 1)
+            
+            nsfw_score = next((item['score'] for item in json_response if item['label'] == 'nsfw'), 0)
+            print(f"NSFW Score: {nsfw_score}")
+            
+            if nsfw_score > 0.1:
+                return None
+                  
+            return self
+            
+        except json.JSONDecodeError as e:
+            raise ImageError(f"NSFW check failed: Invalid response format - {str(e)}")
+        except Exception as e:
+            if attempts > 30:
+                raise ImageError("NSFW check failed after multiple attempts")
+            return self._check_nsfw(attempts + 1)
+    
+    def _convert_color_mode(self):
+        """Handle color mode conversion."""
+        if self.image.mode not in ('RGB', 'RGBA'):
+            self.image = self.image.convert('RGB')
+        elif self.image.mode == 'RGBA':
+            background = Image.new('RGB', self.image.size, (255, 255, 255))
+            background.paste(self.image, mask=self.image.split()[3])
+            self.image = background
+        return self
+    
+    def _resize_for_pixels(self, max_pixels):
+        """Resize image to meet pixel limit."""
+        current_pixels = self.image.width * self.image.height
+        if current_pixels > max_pixels:
+            aspect_ratio = self.image.width / self.image.height
+            if aspect_ratio > 1:
+                new_width = int((max_pixels * aspect_ratio) ** 0.5)
+                new_height = int(new_width / aspect_ratio)
+            else:
+                new_height = int((max_pixels / aspect_ratio) ** 0.5)
+                new_width = int(new_height * aspect_ratio)
+            self.image = self.image.resize((new_width, new_height), Image.LANCZOS)
+        return self
+    
+    def _ensure_dimensions(self, min_size=320, max_size=4096):
+        if (self.image.width < min_size or 
+            self.image.width > max_size or 
+            self.image.height < min_size or 
+            self.image.height > max_size):
+            
+            new_width = min(max(self.image.width, min_size), max_size)
+            new_height = min(max(self.image.height, min_size), max_size)
+            self.image = self.image.resize((new_width, new_height), Image.LANCZOS)
+        
+        return self
+    
+    def encode(self):
+        image_bytes = io.BytesIO()
+        self.image.save(image_bytes, format='PNG', optimize=True)
+        return base64.b64encode(image_bytes.getvalue()).decode('utf8')
+        
+    def process(self, min_size=320, max_size=4096, max_pixels=4194304):
+        """Process image with all necessary transformations."""
+        result = (self
+            ._convert_color_mode()
+            ._resize_for_pixels(max_pixels)
+            ._ensure_dimensions(min_size, max_size)
+            ._check_nsfw())  # Add NSFW check before encoding
+    
+        if result is None:
+            raise ImageError("Image <b>Not Appropriate</b>")
+            
+        return result.encode()
+
+def process_and_encode_image(image, **kwargs):
+    """Process and encode image with default parameters."""
+    try:
+        image = ImageProcessor(image).process(**kwargs)
+        return image
+    except ImageError as e:
+        return str(e)