Update app.py

app.py

@@ -5,6 +5,10 @@ import openai
 import json
 import uuid
 import os
+import base64
+from PIL import Image
+from PIL.PngImagePlugin import PngInfo
+from io import BytesIO
 
 default_urls = ["https://api.anthropic.com", "https://api.openai.com/v1"]
 
@@ -172,6 +176,60 @@ def generate_second_response(endpoint, api_key, model, generated_output, image_m
 
     return response_text
 
+def inject_json_to_png(image, json_data):
+    if isinstance(json_data, str):
+        json_data = json.loads(json_data)
+
+    img = Image.open(image)
+    
+    # Calculate the aspect ratio of the original image
+    width, height = img.size
+    aspect_ratio = width / height
+    
+    # Calculate the cropping dimensions based on the aspect ratio
+    if aspect_ratio > 400 / 600:
+        # Image is wider than 400x600, crop the sides
+        new_width = int(height * 400 / 600)
+        left = (width - new_width) // 2
+        right = left + new_width
+        top = 0
+        bottom = height
+    else:
+        # Image is taller than 400x600, crop the top and bottom
+        new_height = int(width * 600 / 400)
+        left = 0
+        right = width
+        top = (height - new_height) // 2
+        bottom = top + new_height
+    
+    # Perform cropping
+    img = img.crop((left, top, right, bottom))
+    
+    # Resize the cropped image to 400x600 pixels
+    img = img.resize((400, 600), Image.LANCZOS)
+    
+    # Convert the JSON data to bytes
+    json_bytes = json.dumps(json_data).encode('utf-8')
+    
+    # Create a new PNG image with the JSON data injected into the tEXT chunk
+    output = BytesIO()
+    img.save(output, format='PNG')
+    output.seek(0)
+    
+    # Add the tEXT chunk with the tag 'chara'
+    metadata = PngInfo()
+    metadata.add_text("chara", base64.b64encode(json_bytes))
+    
+    # Save the modified PNG image to a BytesIO object
+    output = BytesIO()
+    create_unique_id = str(uuid.uuid4())
+    if json_data['name']:
+        filename = f"{json_data['name']}_{create_unique_id}.png"
+    img_folder = __file__.replace("app.py", f"outputs/")
+    img.save(f"{img_folder}/{filename}", format='PNG', pnginfo=metadata)
+    
+    return f"{img_folder}/{filename}"
+
 # Set up the Gradio interface
 with gr.Blocks() as demo:
     gr.Markdown("# SillyTavern Character Generator")
@@ -179,39 +237,65 @@ with gr.Blocks() as demo:
     #Text explaining that you can use the API key from the Anthropic API or the OpenAI API
     gr.Markdown("You can use the API key from the Anthropic API or the OpenAI API. The API key should start with 'sk-ant-' for Anthropic or 'sk-' for OpenAI.")
     gr.Markdown("Please Note: If you use a proxy it must support the OpenAI or Anthropic standard api calls! khanon does, Openrouter based ones usually do not.")
+    with gr.Tab("JSON Generate"):
+        with gr.Row():
+            with gr.Column():
+                endpoint = gr.Textbox(label="Endpoint", value="https://api.anthropic.com")
+                api_key = gr.Textbox(label="API Key", type="password", placeholder="sk-ant-api03-... or sk-...")
+                model_dropdown = gr.Dropdown(choices=[], label="Select a model")
+                user_prompt = gr.Textbox(label="User Prompt", value="Make me a card for a panther made of translucent pastel colored goo. Its color never changes once it exists but each 'copy' has a different color. The creature comes out of a small jar, seemingly defying physics with its size. It is the size of a real panther, and as strong as one too. By default its female but is able to change gender. It can even split into multiple copies of itself if needed with no change in its own size or mass. Its outside is normally lightly squishy but solid, but on command it can become viscous like non-newtonian fluids. Be descriptive when describing this character, and make sure to describe all of its features in char_persona just like you do in description. Make sure to describe commonly used features in detail (visual, smell, taste, touch, etc).")
+                generate_button = gr.Button("Generate JSON")
 
-    with gr.Row():
-        with gr.Column():
-            endpoint = gr.Textbox(label="Endpoint", value="https://api.anthropic.com")
-            api_key = gr.Textbox(label="API Key", type="password", placeholder="sk-ant-api03-... or sk-...")
-            model_dropdown = gr.Dropdown(choices=[], label="Select a model")
-            user_prompt = gr.Textbox(label="User Prompt", value="Make me a card for a panther made of translucent pastel colored goo. Its color never changes once it exists but each 'copy' has a different color. The creature comes out of a small jar, seemingly defying physics with its size. It is the size of a real panther, and as strong as one too. By default its female but is able to change gender. It can even split into multiple copies of itself if needed with no change in its own size or mass. Its outside is normally lightly squishy but solid, but on command it can become viscous like non-newtonian fluids. Be descriptive when describing this character, and make sure to describe all of its features in char_persona just like you do in description. Make sure to describe commonly used features in detail (visual, smell, taste, touch, etc).")
-            generate_button = gr.Button("Generate JSON")
+            with gr.Column():
+                generated_output = gr.Textbox(label="Generated Output")
+                json_output = gr.Textbox(label="JSON Output")
+                json_download = gr.File(label="Download JSON")
 
-        with gr.Column():
-            generated_output = gr.Textbox(label="Generated Output")
-            json_output = gr.Textbox(label="JSON Output")
-            json_download = gr.File(label="Download JSON")
+        with gr.Row():
+            with gr.Column():
+                image_model = gr.Dropdown(choices=image_prompter, label="Image Model to prompt for", value="SDXL")
+                generate_button_2 = gr.Button("Generate SDXL Prompt")
 
-    with gr.Row():
-        with gr.Column():
-            image_model = gr.Dropdown(choices=image_prompter, label="Image Model to prompt for", value="SDXL")
-            generate_button_2 = gr.Button("Generate SDXL Prompt")
+            with gr.Column():
+                generated_output_2 = gr.Textbox(label="Generated SDXL Prompt")
 
-        with gr.Column():
-            generated_output_2 = gr.Textbox(label="Generated SDXL Prompt")
+        def update_models(api_key):
+            if api_key.startswith("sk-ant-"):
+                return gr.Dropdown(choices=claude_models), gr.Textbox(label="Endpoint", value="https://api.anthropic.com")
+            elif api_key.startswith("sk-"):
+                return gr.Dropdown(choices=openai_models), gr.Textbox(label="Endpoint", value="https://api.openai.com/v1")
+            else:
+                return gr.Dropdown(choices=both_models), gr.Textbox(label="Endpoint", value="https://api.anthropic.com")
 
-    def update_models(api_key):
-        if api_key.startswith("sk-ant-"):
-            return gr.Dropdown(choices=claude_models), gr.Textbox(label="Endpoint", value="https://api.anthropic.com")
-        elif api_key.startswith("sk-"):
-            return gr.Dropdown(choices=openai_models), gr.Textbox(label="Endpoint", value="https://api.openai.com/v1")
-        else:
-            return gr.Dropdown(choices=both_models), gr.Textbox(label="Endpoint", value="https://api.anthropic.com")
+        api_key.change(update_models, inputs=api_key, outputs=[model_dropdown, endpoint])
+
+        generate_button.click(generate_response, inputs=[endpoint, api_key, model_dropdown, user_prompt], outputs=[generated_output, json_output, json_download])
+        generate_button_2.click(generate_second_response, inputs=[endpoint, api_key, model_dropdown, generated_output, image_model], outputs=generated_output_2)
+    with gr.Tab("PNG Inject"):
+        gr.Markdown("# PNG Inject")
+        gr.Markdown("Upload a PNG image and inject JSON content into the PNG. PNG gets resized to 400x600 Center Crop.")
+
+        with gr.Row():
+            with gr.Column():
+                image_input = gr.Image(type="filepath", label="Upload PNG Image")
+                json_input = gr.Textbox(label="JSON Data")
+                json_file_input = gr.File(label="Or Upload JSON File", file_types=[".json"])
+                inject_button = gr.Button("Inject JSON and Download PNG")
+
+            with gr.Column():
+                injected_image_output = gr.File(label="Download Injected PNG")
 
-    api_key.change(update_models, inputs=api_key, outputs=[model_dropdown, endpoint])
+        def inject_json(image, json_data, json_file):
+            if json_file:
+                jsonc = open(json_file,)
+                json_data = json.load(jsonc)
+            if image is None:
+                return None
+            if json_data is None:
+                return None
+            injected_image = inject_json_to_png(image, json_data)
+            return injected_image
 
-    generate_button.click(generate_response, inputs=[endpoint, api_key, model_dropdown, user_prompt], outputs=[generated_output, json_output, json_download])
-    generate_button_2.click(generate_second_response, inputs=[endpoint, api_key, model_dropdown, generated_output, image_model], outputs=generated_output_2)
+        inject_button.click(inject_json, inputs=[image_input, json_input, json_file_input], outputs=injected_image_output)
 
 demo.launch()