Update app.py

app.py

@@ -8,6 +8,7 @@ from PIL import Image
 import requests
 import io
 import base64
+import zipfile
 
 # Set API tokens
 os.environ["REPLICATE_API_TOKEN"] = "r8_Brv0MtpmAiqrXrMrziyUXoSHuFV5hqs1Lw4Mo"
@@ -57,7 +58,6 @@ 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,33 +70,37 @@ def image_to_data_url(image):
     return f"data:image/png;base64,{img_str}"
 
 # Function to inpaint images
-def inpaint_defect(image, prompt):
+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)
-       
-    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
+    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
 
 # Function to generate images from prompts
 def generate_images(prompts):
@@ -121,6 +125,17 @@ 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"):
@@ -162,18 +177,30 @@ 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.Image()
+            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_submit_click_inpaint(image, inpaint_prompt):
-                inpainted_image = inpaint_defect(image, inpaint_prompt)
-                return inpainted_image
+            def on_download_click(images):
+                zip_buffer = download_images_as_zip(images)
+                return zip_buffer
 
             submit_button_inpaint.click(
                 fn=on_submit_click_inpaint,
-                inputs=[image_input, inpaint_prompt_input],
+                inputs=[image_input, inpaint_prompt_input, number_input_inpaint],
                 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()
+    app.launch()