Update app.py

app.py

@@ -8,7 +8,6 @@ from PIL import Image
 import requests
 import io
 import base64
-import zipfile
 
 # Set API tokens
 os.environ["REPLICATE_API_TOKEN"] = "r8_Brv0MtpmAiqrXrMrziyUXoSHuFV5hqs1Lw4Mo"
@@ -58,6 +57,7 @@ def generate_variations(base_prompt, number_of_variations):
         weather = random.choice(weather_conditions)
         
         # Enhance the base prompt with the GPT model
+        #enhanced_prompt = ask_rail_defect_question(base_prompt)
         enhanced_prompt = base_prompt
         full_prompt = f"{enhanced_prompt}, with a {size} defect {location}, observed {weather}."
         variations.append(full_prompt)
@@ -70,37 +70,33 @@ def image_to_data_url(image):
     return f"data:image/png;base64,{img_str}"
 
 # Function to inpaint images
-def inpaint_defect(image, prompt, num_images=1):
+def inpaint_defect(image, prompt):
     if isinstance(image, np.ndarray):
         image = Image.fromarray(image)
     
     image_data_url = image_to_data_url(image)
-    images = []
-
-    for _ in range(num_images):
-        input = {
-            "input_image": image_data_url,
-            "instruction_text": 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="10e63b0e6361eb23a0374f4d9ee145824d9d09f7a31dcd70803193ebc7121430",
-            input=input
-        )
-        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)
-        else:
-            images.append(None)
-    return images
+       
+    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="10e63b0e6361eb23a0374f4d9ee145824d9d09f7a31dcd70803193ebc7121430",
+        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):
@@ -125,17 +121,6 @@ def process_railway_defects(prompt, number_of_images):
     images = generate_images(variations)
     return images
 
-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())
-    zip_buffer.seek(0)
-    return zip_buffer
-
 # UI creation
 with gr.Blocks() as app:
     with gr.Tabs("Prompt Input"):
@@ -177,30 +162,18 @@ with gr.Blocks() as app:
             with gr.Row():
                 image_input = gr.Image(label="Upload Image")
                 inpaint_prompt_input = gr.Textbox(label="Defect Description")
-                number_input_inpaint = gr.Number(label="Number of images to generate", value=1, minimum=1, maximum=10)
                 submit_button_inpaint = gr.Button("Inpaint Defect")
-            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
+            inpainted_image_output = gr.Image()
 
-            def on_download_click(images):
-                zip_buffer = download_images_as_zip(images)
-                return zip_buffer
+            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, number_input_inpaint],
+                inputs=[image_input, inpaint_prompt_input],
                 outputs=inpainted_image_output
             )
 
-            download_button.click(
-                fn=on_download_click,
-                inputs=inpainted_image_output,
-                outputs=gr.File(label="Download Zip")
-            )
-
 if __name__ == "__main__":
     app.launch()