Update app.py

app.py

@@ -11,7 +11,7 @@ import torch
 DEFAULT_BLUR_RADIUS = 2
 DEFAULT_ERODE_ITERATIONS = 1
 DEFAULT_TOLERANCE = 2
-DEFAULT_INFERENCE_STEPS = 30 # Reduced inference steps
+DEFAULT_INFERENCE_STEPS = 20  # Reduced inference steps
 
 # Initialize pipeline with optimizations
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
@@ -62,14 +62,14 @@ def image_to_svg_paths(image, tolerance=DEFAULT_TOLERANCE, blur_radius=DEFAULT_B
         svg_paths.append(path)
     return svg_paths
 
-async def generate_svgs(prompt, color, num_images, fill_color, stroke_color, stroke_width, tolerance, blur_radius, erode_iterations, progress=gr.Progress()):
+async def generate_svgs(prompt, color, fill_color, stroke_color, stroke_width, tolerance, blur_radius, erode_iterations, inference_steps, progress=gr.Progress()):
     """Generate SVGs from a text prompt."""
     color_prompt = f"with {color} colors" if color != "Random" else ""
     full_prompt = f"{prompt}, {color_prompt}"
     svg_outputs = []
-    for i in progress.tqdm(range(num_images)):
+    for i in progress.tqdm(range(3)):  # Fixed to 3 generations
         try:
-            image = pipe(full_prompt, num_inference_steps=DEFAULT_INFERENCE_STEPS, height=256, width=256).images[0].convert("RGB")
+            image = pipe(full_prompt, num_inference_steps=inference_steps, height=256, width=256).images[0].convert("RGB")
             svg_paths = image_to_svg_paths(image, tolerance=tolerance, blur_radius=blur_radius, erode_iterations=erode_iterations)
             svg_string = f'<svg width="256" height="256" xmlns="http://www.w3.org/2000/svg">'
             for path in svg_paths:
@@ -82,7 +82,7 @@ async def generate_svgs(prompt, color, num_images, fill_color, stroke_color, str
             svg_outputs.append(f"Error: {str(e)}")
 
     # Ensure a minimum number of outputs
-    while len(svg_outputs) < 4:
+    while len(svg_outputs) < 3:
         svg_outputs.append("No Image Generated")
     return svg_outputs
 
@@ -92,13 +92,13 @@ with gr.Blocks() as demo:
         with gr.Column():
             prompt_input = gr.Textbox(label="Describe the image", lines=3, placeholder="A cat playing with a ball, vibrant colors")
             color_input = gr.Dropdown(choices=list(COLORS.keys()), label="Desired Colors", value="Random")
-            num_images_input = gr.Slider(minimum=1, maximum=4, value=4, step=1, label="Number of Images")
             fill_color_input = gr.ColorPicker(label="Fill Color", value="#000000")
             stroke_color_input = gr.ColorPicker(label="Stroke Color", value="#000000")
             stroke_width_input = gr.Slider(minimum=1, maximum=5, value=1, step=1, label="Stroke Width")
             tolerance_input = gr.Slider(minimum=1, maximum=10, value=DEFAULT_TOLERANCE, step=1, label="Contour Tolerance")
             blur_radius_input = gr.Slider(minimum=1, maximum=5, value=DEFAULT_BLUR_RADIUS, step=1, label="Blur Radius")
             erode_iterations_input = gr.Slider(minimum=1, maximum=3, value=DEFAULT_ERODE_ITERATIONS, step=1, label="Erode Iterations")
+            inference_steps_input = gr.Slider(minimum=10, maximum=50, value=DEFAULT_INFERENCE_STEPS, step=5, label="Inference Steps")
             generate_button = gr.Button("Generate")
             gr.Markdown("## Prompt Engineering Tips:\n- Be specific and descriptive.\n- Use keywords related to style, color, and composition.\n- Experiment with different prompts to find what works best.")
 
@@ -106,12 +106,11 @@ with gr.Blocks() as demo:
         svg_output1 = gr.HTML(label="SVG 1")
         svg_output2 = gr.HTML(label="SVG 2")
         svg_output3 = gr.HTML(label="SVG 3")
-        svg_output4 = gr.HTML(label="SVG 4")
 
     generate_button.click(
         fn=generate_svgs,
-        inputs=[prompt_input, color_input, num_images_input, fill_color_input, stroke_color_input, stroke_width_input, tolerance_input, blur_radius_input, erode_iterations_input],
-        outputs=[svg_output1, svg_output2, svg_output3, svg_output4],
+        inputs=[prompt_input, color_input, fill_color_input, stroke_color_input, stroke_width_input, tolerance_input, blur_radius_input, erode_iterations_input, inference_steps_input],
+        outputs=[svg_output1, svg_output2, svg_output3],
     )
 
 demo.launch(share=True)