add ngrok endpoint as input

app.py

@@ -19,7 +19,7 @@ class AudioPalette:
         self.image_captioning = ImageCaptioning()
         self.audio_generation = AudioGeneration()
     
-    def generate(self, input_image: PIL.Image.Image):
+    def generate(self, input_image: PIL.Image.Image, ngrok_endpoint: str):
         pace = self.pace_model.predict(input_image)
         print("Pace Prediction Done")
 
@@ -29,7 +29,7 @@ class AudioPalette:
         generated_text = generated_text if generated_text is not None else ""
         prompt = f"Generate a soundtrack for {generated_text} with {pace} beats and the instrument of choice is the guitar, High quality"
 
-        audio_file = self.audio_generation.generate(prompt)
+        audio_file = self.audio_generation.generate(prompt, ngrok_endpoint)
         print("Audio Generation Done")
 
         outputs = [prompt, pace, generated_text, audio_file]
@@ -41,12 +41,23 @@ def main():
 
     demo = gr.Interface(
         fn=model.generate,
-        inputs=gr.Image(
-            type="pil",
-            label="Upload an image",
-            show_label=True,
-            container=True
-        ),
+        inputs=[
+            gr.Image(
+                type="pil",
+                label="Upload an image",
+                show_label=True,
+                container=True
+            ),
+            gr.Textbox(
+                lines=1,
+                placeholder="ngrok endpoint",
+                label="colab endpoint",
+                show_label=True,
+                container=True,
+                type="text",
+                visible=True
+            )
+        ],
         outputs=[
             gr.Textbox(
                 lines=1,

lib/audio_generation.py

@@ -6,13 +6,10 @@ import requests
 
 class AudioGeneration:
     def __init__(self):
-        self.endpoint = os.environ["colab_ngrok_api_endpoint"]
-        self.request_single_endpoint = self.endpoint + "single"
-        self.download_endpoint = self.endpoint + "download"
         self.session = requests.session()
     
     def request_single(self, prompt: str):
-        response = self.session.post(self.request_single_endpoint, json={
+        response = self.session.post(self.endpoint + "single", json={
             "caption": prompt
         })
 
@@ -23,7 +20,7 @@ class AudioGeneration:
         pass
     
     def request_download(self, file_path: str):
-        response = self.session.post(self.download_endpoint, json={
+        response = self.session.post(self.endpoint + "download", json={
             "file_path": file_path
         })
 
@@ -34,7 +31,8 @@ class AudioGeneration:
         
         return audio_file_path
     
-    def generate(self, prompt: typing.Union[str, typing.List[str]]):
+    def generate(self, prompt: typing.Union[str, typing.List[str]], endpoint: str):
+        self.endpoint = endpoint
         if isinstance(prompt, str):
             stored_file_path = self.request_single(prompt)
             audio_file = self.request_download(stored_file_path)