Update app.py

app.py

@@ -22,14 +22,12 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=dtype).to(device)
 
 MAX_SEED = np.iinfo(np.int32).max
-
-MAX_IMAGE_SIZE = 1024
+MAX_IMAGE_SIZE = 512
 
 # Few-shot examples
 few_shot_examples = [
     ("Create a birthday card for friend", "A vibrant birthday card with a colorful confetti background, featuring a large, playful 'Happy Birthday!' in the center. The card has a fun, festive theme with balloons, streamers, and a cupcake with a single lit candle. The message inside reads, 'Wishing you a day full of laughter and joy!'"),
     ("An educational infographic showing the stages of the water cycle with bright, engaging visuals.", "An educational infographic illustrating the water cycle. The diagram shows labeled stages including evaporation, condensation, precipitation, and collection, with arrows guiding the flow. The colors are bright and engaging, with clouds, raindrops, and a sun. The design is simple and clear, suitable for a classroom setting."),
-    
 ]
 
 def generate_detailed_prompt(user_input):
@@ -38,19 +36,18 @@ def generate_detailed_prompt(user_input):
     Given a simple description, create an elaborate and detailed prompt that can be used to generate high-quality images.
     Your response should be concise and no longer than 3 sentences.
     Use the following examples as a guide for the level of detail and creativity expected:
-    
+
     """ + "\n\n".join([f"Input: {input}\nOutput: {output}" for input, output in few_shot_examples]))
 
     human_message = HumanMessage(content=f"Generate a detailed image prompt based on this input, using no more than 3 sentences: {user_input}")
 
     response = llm([system_message, human_message])
     return response.content
-    
+
 @spaces.GPU()
-def generate_image(prompt, seed=0, randomize_seed=False, width=1024, height=1024, num_inference_steps=4):
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-    generator = torch.Generator(device=device).manual_seed(seed if seed is not None else 0)
+def generate_image(prompt, width=512, height=512, num_inference_steps=4):
+    seed = random.randint(0, MAX_SEED)
+    generator = torch.Generator(device=device).manual_seed(seed)
     image = pipe(
         prompt=prompt,
         width=width,
@@ -59,14 +56,10 @@ def generate_image(prompt, seed=0, randomize_seed=False, width=1024, height=1024
         generator=generator,
         guidance_scale=0.0
     ).images[0]
-    
-    return image, seed
 
+    return image
 
 # Gradio UI setup
-import gradio as gr
-
-# Gradio UI
 css = """
 #col-container {
     margin: 0 auto;
@@ -87,13 +80,9 @@ css = """
 #result {
     margin-bottom: 20px;
 }
-#advanced-settings {
-    margin-bottom: 20px;
-}
 """
 
 with gr.Blocks(css=css, theme='gradio/soft') as demo:
-
     with gr.Column(elem_id="col-container"):
         gr.Markdown("""
         # AI-Enhanced Image Generation
@@ -109,61 +98,23 @@ with gr.Blocks(css=css, theme='gradio/soft') as demo:
                 container=False,
                 elem_id="prompt"
             )
-
             run_button = gr.Button("Generate Image", scale=0)
 
         result = gr.Image(label="Result", show_label=False, elem_id="result")
 
-        with gr.Accordion("Advanced Settings", open=False, elem_id="advanced-settings"):
-            seed = gr.Slider(
-                label="Seed",
-                minimum=0,
-                maximum=MAX_SEED,
-                step=1,
-                value=0,
-            )
-
-            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-
-            with gr.Row():
-                width = gr.Slider(
-                    label="Width",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=1024,
-                )
-
-                height = gr.Slider(
-                    label="Height",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=1024,
-                )
-
-            with gr.Row():
-                num_inference_steps = gr.Slider(
-                    label="Number of inference steps",
-                    minimum=1,
-                    maximum=50,
-                    step=1,
-                    value=4,
-                )
-
         gr.Examples(
             examples=[example[0] for example in few_shot_examples],
             inputs=[prompt],
             outputs=[result],
             fn=generate_image,
-            cache_examples=False  # Disable caching to avoid async issues
+            cache_examples=False
         )
 
     gr.on(
         triggers=[run_button.click, prompt.submit],
         fn=generate_image,
-        inputs=[prompt, seed, randomize_seed, width, height, num_inference_steps],
-        outputs=[result, seed]
+        inputs=[prompt],
+        outputs=[result]
     )
 
 demo.launch(share=True)