Update app.py

app.py

@@ -26,7 +26,7 @@ def run_inference(image, model_type, illusion_type, eps_value, num_iterations):
     config = get_inference_configs(eps=eps, n_itr=int(num_iterations))
     
     # Run generative inference
-    output_images, all_steps = model.inference(image, model_type, config)
+    output_image, all_steps = model.inference(image, model_type, config)
     
     # Create animation frames
     frames = []
@@ -35,8 +35,11 @@ def run_inference(image, model_type, illusion_type, eps_value, num_iterations):
         step_pil = Image.fromarray((step_image.permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8))
         frames.append(step_pil)
     
-    # Return the final inferred image and the animation
-    return output_images, gr.Gallery.update(value=frames)
+    # Convert the final output image to PIL
+    final_image = Image.fromarray((output_image.permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8))
+    
+    # Return the final inferred image and the animation frames directly
+    return final_image, frames
 
 # Define the interface
 with gr.Blocks(title="Generative Inference Demo") as demo:
@@ -77,7 +80,7 @@ with gr.Blocks(title="Generative Inference Demo") as demo:
         [os.path.join("stimuli", "Kanizsa_square.jpg"), "robust_resnet50", "Kanizsa", 0.5, 50],
         [os.path.join("stimuli", "face_vase.png"), "robust_resnet50", "Face-Vase", 0.5, 50],
         [os.path.join("stimuli", "figure_ground.png"), "robust_resnet50", "Figure-Ground", 0.7, 100],
-        [os.path.join("stimuli", "NeonColorSaeedi.jpg"), "robust_resnet50", "Neon-Color", 0.3, 80]
+        [os.path.join("stimuli", "Neon_Color_Circle.jpg"), "robust_resnet50", "Neon-Color", 0.3, 80]
     ]
     
     gr.Examples(examples=examples, inputs=[image_input, model_choice, illusion_type, eps_slider, iterations_slider])