Update app.py

app.py

@@ -55,7 +55,7 @@ def generate_image(prompt, width, height, lora_path, trigger_word, steps, custom
 def run_lora(prompt, width, height, lora_path, trigger_word, steps, custom_hash):
     return generate_image(prompt, width, height, lora_path, trigger_word, steps, custom_hash)
 
-# Set up the Gradio interface
+# Gradio interface
 with gr.Blocks() as app:
     gr.Markdown("# LoRA Image Generator")
     
@@ -87,6 +87,16 @@ with gr.Blocks() as app:
         outputs=[output_image, output_hash]
     )
 
+# Gradio API endpoint to serve images based on hash
+@gr.api("/get_image/<hash>")
+def get_image(hash):
+    image_path = f"{hash}.png"
+    if os.path.exists(image_path):
+        # Load and return the image
+        return Image.open(image_path)
+    else:
+        # Return None or a custom message if image not found
+        return "Image not found", 404
+
 if __name__ == "__main__":
-    app.queue()
-    app.launch()
+    app.queue().launch(share=False)