GuardRails

app.py

@@ -44,10 +44,11 @@ with gr.Blocks() as demo:
             """)
             prompt = gr.Textbox(label="Prompt", 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")
             output = gr.Image()
             with gr.Accordion("Advanced Options", open=False):
                 negative_text, width, height, quality, cfg_scale, seed = create_advanced_options()
-            gr.Button("Generate").click(text_to_image, inputs=[prompt, negative_text, height, width, quality, cfg_scale, seed], outputs=output)
+            gr.Button("Generate").click(text_to_image, inputs=[prompt, negative_text, height, width, quality, cfg_scale, seed], outputs=[output, error_box])
 
     with gr.Tab("Inpainting"):
         with gr.Column():
@@ -65,11 +66,12 @@ with gr.Blocks() as demo:
             mask_prompt = gr.Textbox(label="Mask Prompt", placeholder="Describe regions to edit", max_lines=1)
             with gr.Accordion("Mask Image", open=False):
                 mask_image = gr.Image(type='pil', label="Mask Image")
+            error_box = gr.Markdown(visible=False, label="Error")
             output = gr.Image()
             with gr.Accordion("Advanced Options", open=False):
                 negative_text, width, height, quality, cfg_scale, seed = create_advanced_options()
             
-            gr.Button("Generate").click(inpainting, inputs=[image, mask_prompt, mask_image, prompt, negative_text, height, width, quality, cfg_scale, seed], outputs=output)
+            gr.Button("Generate").click(inpainting, inputs=[image, mask_prompt, mask_image, prompt, negative_text, height, width, quality, cfg_scale, seed], outputs=[output, error_box])
 
     with gr.Tab("Outpainting"):
         with gr.Column():
@@ -88,12 +90,13 @@ with gr.Blocks() as demo:
             mask_prompt = gr.Textbox(label="Mask Prompt", placeholder="Describe regions to edit", max_lines=1)
             with gr.Accordion("Mask Image", open=False):
                 mask_image = gr.Image(type='pil', label="Mask Image")
+            error_box = gr.Markdown(visible=False, label="Error")
             output = gr.Image()
             with gr.Accordion("Advanced Options", open=False):
                 outpainting_mode = gr.Radio(choices=["DEFAULT", "PRECISE"], value="DEFAULT", label="Outpainting Mode")
                 negative_text, width, height, quality, cfg_scale, seed = create_advanced_options()
             
-            gr.Button("Generate").click(outpainting, inputs=[image, mask_prompt, mask_image, prompt, negative_text, outpainting_mode, height, width, quality, cfg_scale, seed], outputs=output)
+            gr.Button("Generate").click(outpainting, inputs=[image, mask_prompt, mask_image, prompt, negative_text, outpainting_mode, height, width, quality, cfg_scale, seed], outputs=[output, error_box])
 
     with gr.Tab("Image Variation"):
         with gr.Column():
@@ -106,12 +109,13 @@ with gr.Blocks() as demo:
             with gr.Accordion("Optional Prompt", open=False):
                 prompt = gr.Textbox(label="Prompt", 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")
             output = gr.Image()
             with gr.Accordion("Advanced Options", open=False):
                 similarity_strength = gr.Slider(minimum=0.2, maximum=1.0, step=0.1, value=0.7, label="Similarity Strength")
                 negative_text, width, height, quality, cfg_scale, seed = create_advanced_options()
             
-            gr.Button("Generate").click(image_variation, inputs=[images, prompt, negative_text, similarity_strength, height, width, quality, cfg_scale, seed], outputs=output)
+            gr.Button("Generate").click(image_variation, inputs=[images, prompt, negative_text, similarity_strength, height, width, quality, cfg_scale, seed], outputs=[output, error_box])
 
     with gr.Tab("Image Conditioning"):
         with gr.Column():
@@ -125,12 +129,13 @@ with gr.Blocks() as demo:
             condition_image = gr.Image(type='pil', label="Condition Image")
             prompt = gr.Textbox(label="Prompt", 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")
             output = gr.Image()
             with gr.Accordion("Advanced Options", open=False):
                 control_mode = gr.Radio(choices=["CANNY_EDGE", "SEGMENTATION"], value="CANNY_EDGE", label="Control Mode")
                 control_strength = gr.Slider(minimum=0.0, maximum=1.0, step=0.1, value=0.7, label="Control Strength")
                 negative_text, width, height, quality, cfg_scale, seed = create_advanced_options()
-            gr.Button("Generate").click(image_conditioning, inputs=[condition_image, prompt, negative_text, control_mode, control_strength, height, width, quality, cfg_scale, seed], outputs=output)
+            gr.Button("Generate").click(image_conditioning, inputs=[condition_image, prompt, negative_text, control_mode, control_strength, height, width, quality, cfg_scale, seed], outputs=[output, error_box])
 
     with gr.Tab("Color Guided Content"):
         with gr.Column():
@@ -145,10 +150,11 @@ with gr.Blocks() as demo:
             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)
+            error_box = gr.Markdown(visible=False, label="Error")
             output = gr.Image()
             with gr.Accordion("Advanced Options", open=False):
                 negative_text, width, height, quality, cfg_scale, seed = create_advanced_options()
-            gr.Button("Generate").click(color_guided_content, inputs=[prompt, reference_image, negative_text, colors, height, width, quality, cfg_scale, seed], outputs=output)
+            gr.Button("Generate").click(color_guided_content, inputs=[prompt, reference_image, negative_text, colors, height, width, quality, cfg_scale, seed], outputs=[output, error_box])
 
     with gr.Tab("Background Removal"):
         with gr.Column():
@@ -158,13 +164,30 @@ with gr.Blocks() as demo:
                 </div>
             """)
             image = gr.Image(type='pil', label="Input Image")
+            error_box = gr.Markdown(visible=False, label="Error")
             output = gr.Image()
-            gr.Button("Generate").click(background_removal, inputs=image, outputs=output)
+            gr.Button("Generate").click(background_removal, inputs=image, outputs=[output, error_box])
 
     with gr.Accordion("Tips", open=False):
         gr.Markdown("On Inference Speed: Resolution (width/height), and quality all have an impact on Inference Speed.")
         gr.Markdown("On Negation: For example, consider the prompt \"a rainy city street at night with no people\". The model might interpret \"people\" as a directive of what to include instead of omit. To generate better results, you could use the prompt \"a rainy city street at night\" with a negative prompt \"people\".")
         gr.Markdown("On Prompt Length: When diffusion models were first introduced, they could process only 77 tokens. While new techniques have extended this limit, they remain bound by their training data. AWS Nova Canvas limits input by character length instead, ensuring no characters beyond the set limit are considered in the generated model.")
+
+    gr.Markdown("""
+    <div style="text-align: center;">
+        <h1>Sample Prompts and Results</h1>
+    </div>
+    
+    
+    | Example | Prompt |
+    |:-------:|:-------|
+    | <img src='examples/sample2.png' width='200'> | A whimsical outdoor scene where vibrant flowers and sprawling vines, crafted from an array of colorful fruit leathers and intricately designed candies, flutter with delicate, lifelike butterflies made from translucent, shimmering sweets. Each petal and leaf glistens with a soft, sugary sheen, casting playful reflections. The butterflies, with their candy wings adorned in fruity patterns, flit about, creating a magical, edible landscape that delights the senses. |
+    | <img src='examples/sample4.png' width='200'> | A rugged adventurer's ensemble, crafted for the wild, featuring a khaki jacket adorned with numerous functional pockets, a sun-bleached pith hat with a wide brim, sturdy canvas trousers with reinforced knees, and a pair of weathered leather boots with high-traction soles. Accented with a brass compass pendant and a leather utility belt laden with small tools, the outfit is completed by a pair of aviator sunglasses and a weathered map tucked into a side pocket. |
+    
+    """)
+
+
+
 if __name__ == "__main__":
     demo.launch()
 

functions.py

@@ -1,13 +1,21 @@
 import json
 import io
 import random
+import gradio as gr
 from PIL import Image
 from generate import *
 from typing import Dict, Any
 
 def display_image(image_bytes):
-    image = Image.open(io.BytesIO(image_bytes))
-    return image
+    if isinstance(image_bytes, str):
+        # If we received a string (error message), return it to be displayed
+        return None, gr.update(visible=True, value=image_bytes)
+    elif image_bytes:
+        # If we received image bytes, process and display the image
+        return Image.open(io.BytesIO(image_bytes)), gr.update(visible=False)
+    else:
+        # Handle None case
+        return None, gr.update(visible=False)
 
 def process_optional_params(**kwargs) -> Dict[str, Any]:
     return {k: v for k, v in kwargs.items() if v is not None}
@@ -60,7 +68,9 @@ def text_to_image(prompt, negative_text=None, height=1024, width=1024, quality="
 
 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>")
     # Prepare the inPaintingParams dictionary
     if mask_prompt and mask_image:
         raise ValueError("You must specify either maskPrompt or maskImage, but not both.")
@@ -80,6 +90,9 @@ def inpainting(image, mask_prompt=None, mask_image=None, text=None, negative_tex
 
 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 mask_prompt and mask_image:
         raise ValueError("You must specify either maskPrompt or maskImage, but not both.")
@@ -117,6 +130,9 @@ def image_variation(images, text=None, negative_text=None, similarity_strength=0
 
 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>")
     # Prepare the textToImageParams dictionary
     text_to_image_params = {
         "text": text,
@@ -131,6 +147,10 @@ def image_conditioning(condition_image, text, negative_text=None, control_mode="
 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 not colors:
         colors = "#FF5733,#33FF57,#3357FF,#FF33A1,#33FFF5,#FF8C33,#8C33FF,#33FF8C,#FF3333,#33A1FF"
     
@@ -146,6 +166,10 @@ def color_guided_content(text=None, reference_image=None, negative_text=None, co
 
 def background_removal(image):
     input_image = process_and_encode_image(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>")
+        
     body = json.dumps({
         "taskType": "BACKGROUND_REMOVAL",
         "backgroundRemovalParams": {"image": input_image}

generate.py

@@ -4,6 +4,8 @@ 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
@@ -46,6 +48,8 @@ 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'
 
@@ -59,6 +63,42 @@ class ImageProcessor:
             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."""
+        try:
+            # Save current image temporarily
+            temp_buffer = io.BytesIO()
+            self.image.save(temp_buffer, format='PNG')
+            temp_buffer.seek(0)
+            
+            API_URL = "https://api-inference.huggingface.co/models/Falconsai/nsfw_image_detection"
+            response = requests.request("POST", API_URL, headers=headers, data=temp_buffer.getvalue())
+            decoded_response = response.content.decode("utf-8")
+            
+            json_response = json.loads(decoded_response)
+            
+            if "error" in json_response:
+                time.sleep(json_response["estimated_time"])
+                return self._check_nsfw(attempts+1)
+            
+            scores = {item['label']: item['score'] for item in json_response}
+            nsfw_score = scores.get('nsfw', 0)
+            print(f"NSFW Score: {nsfw_score}")
+            
+            if nsfw_score > 0.1:
+                raise ImageError("Image <b>Not Appropriate</b>")
+                
+            return self
+            
+        except json.JSONDecodeError as e:
+            print(f'JSON Decoding Error: {e}')
+            raise ImageError("NSFW check failed")
+        except Exception as e:
+            print(f'NSFW Check Error: {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'):
@@ -106,6 +146,7 @@ class ImageProcessor:
                 ._convert_color_mode()
                 ._resize_for_pixels(max_pixels)
                 ._ensure_dimensions(min_size, max_size)
+                ._check_nsfw()  # Add NSFW check before encoding
                 .encode())
 
 # Function to generate an image using Amazon Nova Canvas model
@@ -203,18 +244,80 @@ class BedrockClient:
         
         raise ImageError("Unexpected response format.")
 
+def check_rate_limit(body):
+    body = json.loads(body)
+    quality = body.get('imageGenerationConfig', {}).get('quality', 'standard')
+    
+    s3_client = boto3.client(
+        service_name='s3',
+        aws_access_key_id=os.getenv('AWS_ID'),
+        aws_secret_access_key=os.getenv('AWS_SECRET'),
+        region_name=bucket_region
+    )
+    
+    try:
+        # Get current rate limit data
+        response = s3_client.get_object(
+            Bucket=nova_image_bucket,
+            Key='rate-limit/jsonData.json'
+        )
+        rate_data = json.loads(response['Body'].read().decode('utf-8'))
+    except ClientError as e:
+        if e.response['Error']['Code'] == 'NoSuchKey':
+            # Initialize if file doesn't exist
+            rate_data = {'premium': [], 'standard': []}
+        else:
+            raise ImageError(f"Failed to check rate limit: {str(e)}")
+
+    # Get current timestamp
+    current_time = datetime.now().timestamp()
+    # Keep only requests from last minute
+    twenty_minutes_ago = current_time - 1200
+    
+    # Clean up old entries
+    rate_data['premium'] = [t for t in rate_data['premium'] if t > twenty_minutes_ago]
+    rate_data['standard'] = [t for t in rate_data['standard'] if t > twenty_minutes_ago]
+    
+    # Check limits based on quality
+    if quality == 'premium':
+        if len(rate_data['premium']) >= 2:   
+            raise ImageError("<div style='text-align: center;'>Premium rate limit exceeded. Check back later or you use the <a href='https://docs.aws.amazon.com/bedrock/latest/userguide/playgrounds.html'>Bedrock Playground</a>.</div>")
+        rate_data['premium'].append(current_time)
+    else:  # standard
+        if len(rate_data['standard']) >= 4:
+            raise ImageError("<div style='text-align: center;'>Standard rate limit exceeded. Check back later or you use the <a href='https://docs.aws.amazon.com/bedrock/latest/userguide/playgrounds.html'>Bedrock Playground</a>.</div>")
+        rate_data['standard'].append(current_time)
+    
+    # Update rate limit file
+    s3_client.put_object(
+        Bucket=nova_image_bucket,
+        Key='rate-limit/jsonData.json',
+        Body=json.dumps(rate_data),
+        ContentType='application/json'
+    )
+    
+
 def process_and_encode_image(image, **kwargs):
     """Process and encode image with default parameters."""
-    return ImageProcessor(image).process(**kwargs)
+    try:
+        image = ImageProcessor(image).process(**kwargs)
+        return image
+    except ImageError as e:
+        return str(e)
 
 def generate_image(body):
     """Generate image using Bedrock service."""
-    client = BedrockClient(
-        aws_id=os.getenv('AWS_ID'),
-        aws_secret=os.getenv('AWS_SECRET'),
-        model_id='amazon.nova-canvas-v1:0'
-    )
-    return client.generate_image(body)
+    try:
+        check_rate_limit(body)
+        client = BedrockClient(
+            aws_id=os.getenv('AWS_ID'),
+            aws_secret=os.getenv('AWS_SECRET'),
+            model_id='amazon.nova-canvas-v1:0'
+        )
+        return client.generate_image(body)
+    except ImageError as e:
+        return str(e)
+        
 
 def generate_prompt(body):
     client = BedrockClient(

requirements.txt

@@ -1,4 +1,5 @@
 Pillow
 boto3
 python-dotenv
-gradio
+gradio
+requests