Update app.py

app.py

@@ -5,6 +5,7 @@ from PIL import Image
 from io import BytesIO
 import numpy as np
 import os
+import random
 
 API_URL = 'https://hub.societyai.com/models/flux-1-schnell/infer'
 API_TOKEN = os.environ.get("SAI_API_TOKEN", "")
@@ -19,22 +20,27 @@ with gr.Blocks(css="footer {visibility: hidden}") as demo:
     theme = gr.Dropdown(label="Theme", choices=["Cyberpunk", "Fantasy", "Anime", "Dreamscape"], value="Cyberpunk")
     color = gr.Dropdown(label="Color", choices=["Pink", "Green", "Blue", "Red"], value="Pink")
 
+    # Checkbox for randomize seed
+    randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
+
     # Fixed parameters
-    width = 512
-    height = 512
+    width = 256
+    height = 256
     num_steps = 4
     guidance_scale = 7.5
-    seed = 123
-    randomize_seed = False
+    fixed_seed = 123
 
     generate_button = gr.Button("Generate Image")
     output_image = gr.Image(type="numpy", label="Generated Image")
     message = gr.Textbox(label="Status", interactive=False)
 
-    def generate_image(avatar, hair, theme, color):
+    def generate_image(avatar, hair, theme, color, randomize_seed):
         # Construct the prompt
         prompt = f"image of a {avatar} with {hair} hair, in a {theme} style with {color} as the main color"
 
+        # Seed logic
+        seed = random.randint(0, MAX_SEED) if randomize_seed else fixed_seed
+
         # Validation: Ensure width and height are divisible by 8
         if width % 8 != 0 or height % 8 != 0:
             return None, "Error: Both width and height must be divisible by 8."
@@ -127,7 +133,7 @@ with gr.Blocks(css="footer {visibility: hidden}") as demo:
 
     generate_button.click(
         generate_image,
-        inputs=[avatar, hair, theme, color],
+        inputs=[avatar, hair, theme, color, randomize_seed],
         outputs=[output_image, message]
     )