Update app.py

app.py

@@ -4,7 +4,7 @@ import io
 from PIL import Image
 from together import Together
 
-def generate_image(api_key, prompt):
+def generate_image(api_key, prompt, width, height, steps):
     try:
         # Initialize the Together AI client with the provided API key
         client = Together(api_key=api_key)
@@ -13,9 +13,9 @@ def generate_image(api_key, prompt):
         response = client.images.generate(
             prompt=prompt,
             model="black-forest-labs/FLUX.1-schnell-Free",
-            width=1024,
-            height=768,
-            steps=4,
+            width=width,
+            height=height,
+            steps=steps,
             n=1,
             response_format="b64_json"
         )
@@ -37,7 +37,7 @@ with gr.Blocks() as app:
     gr.Markdown("Generate images using the FLUX.1-schnell-Free model from Together AI")
     
     with gr.Row():
-        with gr.Column():
+        with gr.Column(scale=1):
             api_key_input = gr.Textbox(
                 label="Together API Key",
                 placeholder="Enter your Together API key here...",
@@ -48,23 +48,50 @@ with gr.Blocks() as app:
                 placeholder="A beautiful sunset over mountains...",
                 lines=3
             )
+            
+            with gr.Row():
+                width_input = gr.Number(
+                    label="Width",
+                    value=1024,
+                    min=256,
+                    max=2048,
+                    step=64
+                )
+                height_input = gr.Number(
+                    label="Height",
+                    value=768,
+                    min=256,
+                    max=2048,
+                    step=64
+                )
+                steps_input = gr.Slider(
+                    label="Steps",
+                    minimum=1,
+                    maximum=50,
+                    value=4,
+                    step=1
+                )
+                
             generate_button = gr.Button("Generate Image")
         
-        with gr.Column():
+        with gr.Column(scale=1):
             image_output = gr.Image(label="Generated Image")
     
     generate_button.click(
         fn=generate_image,
-        inputs=[api_key_input, prompt_input],
+        inputs=[api_key_input, prompt_input, width_input, height_input, steps_input],
         outputs=image_output
     )
     
     gr.Markdown("""
     ## Instructions
-    1. Enter your Together AI API key (get one from https://www.together.ai)
+    1. Enter your Together API Key (get one from https://www.together.ai)
     2. Enter a descriptive prompt in the text box
-    3. Click "Generate Image" 
-    4. Wait for the image to be generated
+    3. Adjust the width, height, and steps as needed
+        - Width and height control the dimensions of the generated image
+        - Steps control the number of diffusion steps (higher = more detail but slower)
+    4. Click "Generate Image" 
+    5. Wait for the image to be generated
     
     Note: Your API key is not stored and is only used for the current session.
     """)