Update app.py

app.py

@@ -1,9 +1,13 @@
 import gradio as gr
 import replicate
 import os
-from huggingface_hub import InferenceClient
 import random
 import openai
+import numpy as np
+from PIL import Image
+import requests
+import io
+import base64
 
 # Set API tokens
 os.environ["REPLICATE_API_TOKEN"] = "r8_Brv0MtpmAiqrXrMrziyUXoSHuFV5hqs1Lw4Mo"
@@ -13,8 +17,6 @@ rep_client = replicate.Client()
 # Set your OpenAI API key
 OPENAI_API_KEY = "sk-proj-5iy4bwrqAW8GpguiEawaT3BlbkFJ8p88lLSjOCeDbxWsAOlr"
 openai.api_key = OPENAI_API_KEY
-# Initialize the Replicate client
-rep_client = replicate.Client()
 
 # Predefined prompts for the dropdown
 predefined_prompts = [
@@ -26,9 +28,7 @@ predefined_prompts = [
     "Damaged railway platform"
 ]
 
-
 def ask_rail_defect_question(question, model_name='ft:gpt-3.5-turbo-0125:personal::99NsSAeQ'):
-    openai.api_key = OPENAI_API_KEY
     response = openai.ChatCompletion.create(
         model=model_name,
         messages=[
@@ -63,6 +63,41 @@ def generate_variations(base_prompt, number_of_variations):
         variations.append(full_prompt)
     return variations
 
+def image_to_data_url(image):
+    buffered = io.BytesIO()
+    image.save(buffered, format="PNG")
+    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):
+    if isinstance(image, np.ndarray):
+        image = Image.fromarray(image)
+    
+    image_data_url = image_to_data_url(image)
+       
+    input = {
+      "image": image_data_url,
+      "prompt": prompt,
+      "scheduler": "K_EULER_ANCESTRAL",
+      "num_outputs": 1,
+      "guidance_scale": 7.5,
+      "num_inference_steps": 100,
+      "image_guidance_scale": 1.5
+    }
+
+    prediction = rep_client.predictions.create(
+        version="30c1d0b916a6f8efce20493f5d61ee27491ab2a60437c13c588468b9810ec23f",
+        input = input
+    )
+    prediction.wait()
+    if prediction.status == "succeeded":
+        image_url = prediction.output[0]
+        response = requests.get(image_url)
+        image = Image.open(io.BytesIO(response.content))
+        return image
+    return None
+
 # Function to generate images from prompts
 def generate_images(prompts):
     images = []
@@ -85,8 +120,6 @@ def process_railway_defects(prompt, number_of_images):
     variations = generate_variations(prompt, number_of_images)
     images = generate_images(variations)
     return images
-    
-
 
 # UI creation
 with gr.Blocks() as app:
@@ -102,8 +135,6 @@ with gr.Blocks() as app:
                 images = process_railway_defects(prompt, number_of_images)
                 return images
 
-
-
             submit_button_dropdown.click(
                 fn=on_submit_click_dropdown,
                 inputs=[prompt_input, number_input_dropdown],
@@ -118,8 +149,8 @@ with gr.Blocks() as app:
             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
+                images = process_railway_defects(custom_prompt, number_of_images)
+                return images
 
             submit_button_custom.click(
                 fn=on_submit_click_custom,
@@ -127,5 +158,22 @@ with gr.Blocks() as app:
                 outputs=image_outputs_custom
             )
         
+        with gr.Tab("Inpaint Defect"):
+            with gr.Row():
+                image_input = gr.Image(label="Upload Image")
+                inpaint_prompt_input = gr.Textbox(label="Defect Description")
+                submit_button_inpaint = gr.Button("Inpaint Defect")
+            inpainted_image_output = gr.Image()
+
+            def on_submit_click_inpaint(image, inpaint_prompt):
+                inpainted_image = inpaint_defect(image, inpaint_prompt)
+                return inpainted_image
+
+            submit_button_inpaint.click(
+                fn=on_submit_click_inpaint,
+                inputs=[image_input, inpaint_prompt_input],
+                outputs=inpainted_image_output
+            )
+
 if __name__ == "__main__":
-    app.launch()
+    app.launch()