Update app.py

app.py

@@ -90,7 +90,7 @@ example_images = [
 
 
 # Main function to process the uploaded image
-def process_image(img):
+def process_image(img, generate_image=True):
     print("Starting prediction...")
     predicted_class, _, probs = learn.predict(img)
     print(f"Prediction complete: {predicted_class}")
@@ -109,16 +109,19 @@ def process_image(img):
     try:
         client = OpenAI()
 
-        result = client.images.generate(
-            model="gpt-image-1",
-            prompt=random.choice(prompt_templates).format(flower=predicted_class),
-            size="1024x1024",
-            background="transparent",
-            quality="medium"
-        )
+        if generate_image:
+            result = client.images.generate(
+                model="gpt-image-1",
+                prompt=random.choice(prompt_templates).format(flower=predicted_class),
+                size="1024x1024",
+                background="transparent",
+                quality="medium"
+            )
 
-        image_base64 = result.data[0].b64_json
-        generated_image = base64.b64decode(image_base64)
+            image_base64 = result.data[0].b64_json
+            generated_image = base64.b64decode(image_base64)
+        else:
+            generated_image = None
 
     except Exception as e:
         print(f"Error generating image: {e}")
@@ -160,12 +163,12 @@ with gr.Blocks() as demo:
         examples=example_images,
         inputs=input_image,
         examples_per_page=6,
-        fn=process_image,
+        fn=lambda img: process_image(img, generate_image=False),
         outputs=[label_output, generated_image, wiki_output, endangerment_output]
     )
 
     input_image.change(
-        fn=process_image,
+        fn=lambda img: process_image(img, generate_image=True),
         inputs=input_image,
         outputs=[label_output, generated_image, wiki_output, endangerment_output]
     )