Update app.py

app.py

@@ -12,12 +12,13 @@ import zipfile
 
 # Set API tokens
 os.environ["REPLICATE_API_TOKEN"] = "r8_Brv0MtpmAiqrXrMrziyUXoSHuFV5hqs1Lw4Mo"
-os.environ["OPENAI_API_KEY"] = "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"
-openai.api_key = os.getenv("OPENAI_API_KEY")
-
 # Initialize the Replicate client
 rep_client = replicate.Client()
 
+# Set your OpenAI API key
+OPENAI_API_KEY = "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"
+openai.api_key = OPENAI_API_KEY
+
 # Predefined prompts for the dropdown
 predefined_prompts = [
     "Missing bolts on railway track",
@@ -44,6 +45,7 @@ def ask_rail_defect_question(question, model_name='ft:gpt-3.5-turbo-0125:persona
     )
     return response.choices[0].message['content']
 
+# Function to generate variations enhanced by the GPT model
 def generate_variations(base_prompt, number_of_variations):
     locations = ["on the left side", "on the right side", "at the top", "at the bottom", "in the center"]
     sizes = ["small", "medium", "large", "tiny", "huge"]
@@ -55,8 +57,10 @@ def generate_variations(base_prompt, number_of_variations):
         size = random.choice(sizes)
         weather = random.choice(weather_conditions)
         
-        enhanced_prompt = f"{base_prompt}, with a {size} defect {location}, observed {weather}."
-        variations.append(enhanced_prompt)
+        # Enhance the base prompt with the GPT model
+        enhanced_prompt = base_prompt
+        full_prompt = f"{enhanced_prompt}, with a {size} defect {location}, observed {weather}."
+        variations.append(full_prompt)
     return variations
 
 def image_to_data_url(image):
@@ -65,9 +69,11 @@ def image_to_data_url(image):
     img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
     return f"data:image/png;base64,{img_str}"
 
+# Function to inpaint images
 def inpaint_defect(image, prompt, num_images=1):
     if isinstance(image, np.ndarray):
         image = Image.fromarray(image)
+    
     image_data_url = image_to_data_url(image)
     images = []
 
@@ -81,6 +87,7 @@ def inpaint_defect(image, prompt, num_images=1):
             "num_inference_steps": 100,
             "image_guidance_scale": 1.5
         }
+
         prediction = rep_client.predictions.create(
             version="10e63b0e6361eb23a0374f4d9ee145824d9d09f7a31dcd70803193ebc7121430",
             input=input
@@ -91,8 +98,11 @@ def inpaint_defect(image, prompt, num_images=1):
             response = requests.get(image_url)
             img = Image.open(io.BytesIO(response.content))
             images.append(img)
+        else:
+            images.append(None)
     return images
 
+# Function to generate images from prompts
 def generate_images(prompts):
     images = []
     for prompt in prompts:
@@ -104,6 +114,8 @@ def generate_images(prompts):
             prediction.wait()
             if prediction.status == "succeeded" and prediction.output:
                 images.append(prediction.output[0])
+            else:
+                images.append("Failed to generate image.")
         except Exception as e:
             images.append(f"Error: {str(e)}")
     return images
@@ -134,8 +146,12 @@ with gr.Blocks() as app:
                 submit_button_dropdown = gr.Button("Generate")
             image_outputs_dropdown = gr.Gallery()
 
+            def on_submit_click_dropdown(prompt, number_of_images):
+                images = process_railway_defects(prompt, number_of_images)
+                return images
+
             submit_button_dropdown.click(
-                fn=process_railway_defects,
+                fn=on_submit_click_dropdown,
                 inputs=[prompt_input, number_input_dropdown],
                 outputs=image_outputs_dropdown
             )
@@ -147,8 +163,12 @@ with gr.Blocks() as app:
                 submit_button_custom = gr.Button("Generate")
             image_outputs_custom = gr.Gallery()
 
+            def on_submit_click_custom(custom_prompt, number_of_images):
+                images = process_railway_defects(custom_prompt, number_of_images)
+                return images
+
             submit_button_custom.click(
-                fn=process_railway_defects,
+                fn=on_submit_click_custom,
                 inputs=[custom_prompt_input, number_input_custom],
                 outputs=image_outputs_custom
             )
@@ -162,48 +182,25 @@ with gr.Blocks() as app:
             inpainted_image_output = gr.Gallery()
             download_button = gr.Button("Download Images as Zip")
 
+            def on_submit_click_inpaint(image, inpaint_prompt, number_of_images):
+                inpainted_images = inpaint_defect(image, inpaint_prompt, num_images=number_of_images)
+                return inpainted_images
+
+            def on_download_click(images):
+                zip_buffer = download_images_as_zip(images)
+                return zip_buffer
+
             submit_button_inpaint.click(
-                fn=inpaint_defect,
+                fn=on_submit_click_inpaint,
                 inputs=[image_input, inpaint_prompt_input, number_input_inpaint],
                 outputs=inpainted_image_output
             )
 
             download_button.click(
-                fn=download_images_as_zip,
+                fn=on_download_click,
                 inputs=inpainted_image_output,
                 outputs=gr.File(label="Download Zip")
             )
-        
-        with gr.Tab("Feedback"):
-            feedback_input = gr.Textbox(label="Your feedback")
-            like_button = gr.Button("Like")
-            dislike_button = gr.Button("Dislike")
-            feedback_response = gr.Text(label="Feedback Response")
-
-            def handle_feedback(feedback):
-                sentiment_response = ask_rail_defect_question(feedback, model_name='ft:gpt-3.5-turbo-0125:personal::99NsSAeQ')
-                if "negative" in sentiment_response:
-                    return "Sorry for your experience, and thanks for your feedback. We are updating our system to generate more satisfactory samples."
-                else:
-                    return "Thanks for your positive feedback!"
-
-            feedback_input.change(
-                fn=handle_feedback,
-                inputs=feedback_input,
-                outputs=feedback_response
-            )
-
-            like_button.click(
-                fn=lambda: "Thanks for your positive feedback!",
-                inputs=None,
-                outputs=feedback_response
-            )
-
-            dislike_button.click(
-                fn=lambda: "Sorry for your experience, and thanks for your feedback. We are updating our system to generate more satisfactory samples.",
-                inputs=None,
-                outputs=feedback_response
-            )
 
 if __name__ == "__main__":
     app.launch()