Update app.py

app.py

@@ -1,44 +1,64 @@
 import gradio as gr
 from main import main
 from arguments import parse_args
+import os
 
-def generate_image(prompt):
+def generate_image(prompt, model, num_iterations, learning_rate, progress = gr.Progress(track_tqdm=True)):
     # Set up arguments
     args = parse_args()
     args.task = "single"
     args.prompt = prompt
-    args.model = "sd-turbo"  # or another supported model
+    args.model = model
+    args.n_iters = num_iterations
+    args.lr = learning_rate
     args.cache_dir = "./HF_model_cache"
     args.save_dir = "./outputs"
     args.save_all_images = True
     
-    # Run the main function
-    main(args)
+    try:
+        # Run the main function
+        main(args)
 
-    settings = (
-        f"{args.model}{'_' + args.prompt if args.task == 't2i-compbench' else ''}"
-        f"{'_no-optim' if args.no_optim else ''}_{args.seed if args.task != 'geneval' else ''}"
-        f"_lr{args.lr}_gc{args.grad_clip}_iter{args.n_iters}"
-        f"_reg{args.reg_weight if args.enable_reg else '0'}"
-        f"{'_pickscore' + str(args.pickscore_weighting) if args.enable_pickscore else ''}"
-        f"{'_clip' + str(args.clip_weighting) if args.enable_clip else ''}"
-        f"{'_hps' + str(args.hps_weighting) if args.enable_hps else ''}"
-        f"{'_imagereward' + str(args.imagereward_weighting) if args.enable_imagereward else ''}"
-        f"{'_aesthetic' + str(args.aesthetic_weighting) if args.enable_aesthetic else ''}"
-    )
+        settings = (
+            f"{args.model}{'_' + args.prompt if args.task == 't2i-compbench' else ''}"
+            f"{'_no-optim' if args.no_optim else ''}_{args.seed if args.task != 'geneval' else ''}"
+            f"_lr{args.lr}_gc{args.grad_clip}_iter{args.n_iters}"
+            f"_reg{args.reg_weight if args.enable_reg else '0'}"
+            f"{'_pickscore' + str(args.pickscore_weighting) if args.enable_pickscore else ''}"
+            f"{'_clip' + str(args.clip_weighting) if args.enable_clip else ''}"
+            f"{'_hps' + str(args.hps_weighting) if args.enable_hps else ''}"
+            f"{'_imagereward' + str(args.imagereward_weighting) if args.enable_imagereward else ''}"
+            f"{'_aesthetic' + str(args.aesthetic_weighting) if args.enable_aesthetic else ''}"
+        )
 
-    save_dir = f"{args.save_dir}/{args.task}/{settings}/{args.prompt}"
+        save_dir = f"{args.save_dir}/{args.task}/{settings}/{args.prompt}"
+        
+        # Return the path to the generated image
+        image_path = f"{save_dir}/best_image.png"
+        
+        if os.path.exists(image_path):
+            return image_path, f"Image generated successfully and saved at {image_path}"
+        else:
+            return None, "Image generation completed, but the file was not found."
     
-    # Return the path to the generated image
-    return f"{save_dir}/best_image.png"
+    except Exception as e:
+        return None, f"An error occurred: {str(e)}"
 
 # Create Gradio interface
 iface = gr.Interface(
     fn=generate_image,
-    inputs="text",
-    outputs="image",
+    inputs=[
+        gr.Textbox(label="Prompt"),
+        gr.Dropdown(["sd-turbo", "sdxl-turbo", "pixart", "hyper-sd"], label="Model"),
+        gr.Slider(minimum=10, maximum=100, step=10, value=50, label="Number of Iterations"),
+        gr.Slider(minimum=0.1, maximum=10.0, step=0.1, value=5.0, label="Learning Rate"),
+    ],
+    outputs=[
+        gr.Image(type="filepath", label="Generated Image"),
+        gr.Textbox(label="Status")
+    ],
     title="ReNO Image Generation",
-    description="Enter a prompt to generate an image using ReNO."
+    description="Enter a prompt to generate an image using ReNO. Adjust the model and parameters as needed."
 )
 
 # Launch the app