Update app.py

app.py

@@ -1,68 +1,11 @@
-import gradio as gr
-import replicate
-import os
-import random
-import openai
-import numpy as np
-from PIL import Image
-import requests
 import io
 import base64
+import requests
+import numpy as np
+from PIL import Image
+import gradio as gr
 import zipfile
 
-# Set API tokens
-os.environ["REPLICATE_API_TOKEN"] = "r8_Brv0MtpmAiqrXrMrziyUXoSHuFV5hqs1Lw4Mo"
-# 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",
-    "Cracks on railway track",
-    "Overgrown vegetation near railway track",
-    "Broken railings on railway bridge",
-    "Debris on railway track",
-    "Damaged railway platform"
-]
-
-def ask_rail_defect_question(question, model_name='ft:gpt-3.5-turbo-0125:personal::99NsSAeQ'):
-    response = openai.ChatCompletion.create(
-        model=model_name,
-        messages=[
-            {
-                "role": "system",
-                "content": "The assistant is knowledgeable about rail defects and can answer questions related to them.",
-            },
-            {
-                "role": "user",
-                "content": question,
-            }
-        ],
-    )
-    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"]
-    weather_conditions = ["under cold conditions", "during hot weather", "in dry weather", "in humid conditions", "under varying temperatures"]
-
-    variations = []
-    for _ in range(number_of_variations):
-        location = random.choice(locations)
-        size = random.choice(sizes)
-        weather = random.choice(weather_conditions)
-        
-        # 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):
     buffered = io.BytesIO()
     image.save(buffered, format="PNG")
@@ -78,7 +21,7 @@ def inpaint_defect(image, prompt, num_images=1):
     images = []
 
     for _ in range(num_images):
-        input = {
+        input_data = {
             "input_image": image_data_url,
             "instruction_text": prompt,
             "scheduler": "K_EULER_ANCESTRAL",
@@ -90,16 +33,25 @@ def inpaint_defect(image, prompt, num_images=1):
 
         prediction = rep_client.predictions.create(
             version="10e63b0e6361eb23a0374f4d9ee145824d9d09f7a31dcd70803193ebc7121430",
-            input=input
+            input=input_data
         )
         prediction.wait()
+
         if prediction.status == "succeeded":
             image_url = prediction.output[0]
-            response = requests.get(image_url)
-            img = Image.open(io.BytesIO(response.content))
-            images.append(img)
+            print(f"Generated Image URL: {image_url}")  # Debugging line to check the URL
+            try:
+                response = requests.get(image_url)
+                response.raise_for_status()  # Check for HTTP errors
+                img = Image.open(io.BytesIO(response.content))
+                images.append(img)
+            except requests.exceptions.RequestException as e:
+                print(f"Error fetching image from URL: {e}")
+                images.append(None)
         else:
+            print(f"Prediction failed: {prediction.status}")
             images.append(None)
+    
     return images
 
 # Function to generate images from prompts
@@ -129,10 +81,11 @@ def download_images_as_zip(images):
     zip_buffer = io.BytesIO()
     with zipfile.ZipFile(zip_buffer, 'w') as zf:
         for i, img in enumerate(images):
-            img_buffer = io.BytesIO()
-            img.save(img_buffer, format='PNG')
-            img_buffer.seek(0)
-            zf.writestr(f'image_{i + 1}.png', img_buffer.read())
+            if img is not None:
+                img_buffer = io.BytesIO()
+                img.save(img_buffer, format='PNG')
+                img_buffer.seek(0)
+                zf.writestr(f'image_{i + 1}.png', img_buffer.read())
     zip_buffer.seek(0)
     return zip_buffer
 
@@ -202,5 +155,4 @@ with gr.Blocks() as app:
                 outputs=gr.File(label="Download Zip")
             )
 
-if __name__ == "__main__":
-    app.launch()
+app.launch()