Reorganize structure for less code clutter

app.py

@@ -1,109 +1,63 @@
 from dotenv import load_dotenv
+from functools import partial
 import gradio as gr
 from huggingface_hub import InferenceClient
-from io import BytesIO
-from os import path, unlink
-import pandas as pd
-from pandas import DataFrame
-from PIL.Image import Image, open as open_image
-import requests
-import tempfile
-
-REQUEST_TIMEOUT = 45
-
-TEXT_TO_IMAGE_MODEL = "black-forest-labs/FLUX.1-dev"
-IMAGE_CLASSIFICATION_MODEL = "prithivMLmods/Trash-Net"
-
-# Load environment variables from .env file
-load_dotenv()
-
-client = InferenceClient()
-
-def save_image_to_temp_file(image: Image) -> str:
-    image_format = image.format if image.format else 'PNG'
-    format_extension = image_format.lower() if image_format else 'png'
-    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=f".{format_extension}")
-    temp_path = temp_file.name
-    temp_file.close()
-    image.save(temp_path, format=image_format)
-    return temp_path
-
-def text_to_image(prompt: str) -> Image:
-    return client.text_to_image(prompt, model=TEXT_TO_IMAGE_MODEL)
+from image_classification import image_classification
+from text_to_image import text_to_image
+
+
+class App:
+
+    def __init__(self, client: InferenceClient):
+        self.client = client
+
+    def run(self):
+        with gr.Blocks(title="AI Building Blocks") as demo:
+            gr.Markdown("# AI Building Blocks")
+            gr.Markdown("A gallery of building blocks for building AI applications")
+            with gr.Tabs():
+                with gr.Tab("Text-to-image Generation"):
+                    gr.Markdown("Generate an image from a text prompt.")
+                    text_to_image_prompt = gr.Textbox(label="Prompt", value="A panda under a giant mushroom next to a pumpkin")
+                    text_to_image_generate_button = gr.Button("Generate")
+                    text_to_image_output = gr.Image(label="Image", type="pil")
+                    text_to_image_generate_button.click(
+                        fn=partial(text_to_image, self.client),
+                        inputs=text_to_image_prompt,
+                        outputs=text_to_image_output
+                    )
+                with gr.Tab("Image Classification"):
+                    gr.Markdown("Classify a recyclable item as one of: cardboard, glass, metal, paper, plastic, or other using [Trash-Net](https://huggingface.co/prithivMLmods/Trash-Net).")
+                    with gr.Row():
+                        with gr.Column():
+                            image_classification_url_input = gr.Textbox(
+                                label="Image URL",
+                                value="https://campuslifeservices.ucsf.edu/upload/facilities/galleries/cardboard_0.jpg",
+                                placeholder="Enter the URL of the image to classify",
+                                scale=2
+                            )
+                            image_classification_image_preview = gr.Image(label="Image Preview", type="pil")
+                        image_classification_upload_input = gr.Image(
+                            label="Or Upload Image",
+                            type="pil",
+                            scale=2
+                        )
+                    image_classification_button = gr.Button("Classify")
+                    image_classification_output = gr.Dataframe(
+                        label="Classification Results",
+                        headers=["Label", "Probability"],
+                        interactive=False
+                    )
+                    image_classification_button.click(
+                        fn=partial(image_classification, self.client),
+                        inputs=[image_classification_url_input, image_classification_upload_input],
+                        outputs=[image_classification_image_preview, image_classification_output]
+                    )
 
-def image_classification(image_url: str | None, uploaded_image: Image | None) -> tuple[Image | None, DataFrame]:
-    temp_file_path = None
-    try:
-        if uploaded_image is not None and image_url and image_url.strip():
-            raise gr.Error("Both an image URL and an uploaded image were provided. Please provide only one or the other.")
-        elif uploaded_image is not None:
-            temp_file_path = save_image_to_temp_file(uploaded_image)
-            classifications = client.image_classification(temp_file_path, model=IMAGE_CLASSIFICATION_MODEL)
-            image = None
-        elif image_url and image_url.strip():
-            try:
-                response = requests.get(image_url, timeout=REQUEST_TIMEOUT)
-                response.raise_for_status()
-                image = open_image(BytesIO(response.content))
-                temp_file_path = save_image_to_temp_file(image)
-                classifications = client.image_classification(temp_file_path, model=IMAGE_CLASSIFICATION_MODEL)
-            except Exception as e:
-                raise gr.Error(f"Failed to fetch image from URL: {str(e)}")
-        else:
-            raise gr.Error("Please either provide an image URL or upload an image.")
-        df = pd.DataFrame([
-            {"Label": classification.label, "Probability": f"{classification.score:.2%}"}
-            for classification in classifications
-        ])
-        return image, df
-    finally:
-        # Clean up temporary file.
-        if temp_file_path and path.exists(temp_file_path):
-            try:
-                unlink(temp_file_path)
-            except Exception:
-                pass # Ignore clean-up errors.
+            demo.launch()
 
-with gr.Blocks(title="AI Building Blocks") as demo:
-    gr.Markdown("# AI Building Blocks")
-    gr.Markdown("A gallery of building blocks for building AI applications")
-    with gr.Tabs():
-        with gr.Tab("Text-to-image Generation"):
-            gr.Markdown("Generate an image from a text prompt.")
-            text_to_image_prompt = gr.Textbox(label="Prompt", value="A panda under a giant mushroom next to a pumpkin")
-            text_to_image_generate_button = gr.Button("Generate")
-            text_to_image_output = gr.Image(label="Image", type="pil")
-            text_to_image_generate_button.click(
-                fn=text_to_image,
-                inputs=text_to_image_prompt,
-                outputs=text_to_image_output
-            )
-        with gr.Tab("Image Classification"):
-            gr.Markdown("Classify a recyclable item as one of: cardboard, glass, metal, paper, plastic, or other using [Trash-Net](https://huggingface.co/prithivMLmods/Trash-Net).")
-            with gr.Row():
-                with gr.Column():
-                    image_classification_url_input = gr.Textbox(
-                        label="Image URL",
-                        value="https://campuslifeservices.ucsf.edu/upload/facilities/galleries/cardboard_0.jpg",
-                        placeholder="Enter the URL of the image to classify",
-                        scale=2
-                    )
-                    image_classification_image_preview = gr.Image(label="Image Preview", type="pil")
-                image_classification_upload_input = gr.Image(
-                    label="Or Upload Image",
-                    type="pil",
-                    scale=2
-                )
-            image_classification_button = gr.Button("Classify")
-            image_classification_output = gr.Dataframe(
-                label="Classification Results",
-                headers=["Label", "Probability"],
-                interactive=False
-            )
-            image_classification_button.click(
-                fn=image_classification,
-                inputs=[image_classification_url_input, image_classification_upload_input],
-                outputs=[image_classification_image_preview, image_classification_output]
-            )
 
-    demo.launch()
+if __name__ == "__main__":
+    load_dotenv()
+    app = App(InferenceClient())
+    app.run()

image_classification.py

@@ -0,0 +1,45 @@
+import gradio as gr
+from huggingface_hub import InferenceClient
+from io import BytesIO
+from os import path, unlink, getenv
+from PIL.Image import Image, open as open_image
+import pandas as pd
+from pandas import DataFrame
+import requests
+from utils import save_image_to_temp_file
+
+
+def image_classification(client: InferenceClient, image_url: str | None, image: Image | None) -> tuple[Image | None, DataFrame]:
+    temp_file_path = None
+    try:
+        if image is not None and image_url and image_url.strip():
+            raise gr.Error("Both an image URL and an uploaded image were provided. Please provide only one or the other.")
+        elif image is not None:
+            temp_file_path = save_image_to_temp_file(image)
+            classifications = client.image_classification(temp_file_path, model=getenv("IMAGE_CLASSIFICATION_MODEL"))
+            image = None
+        elif image_url and image_url.strip():
+            try:
+                response = requests.get(image_url, timeout=int(getenv("REQUEST_TIMEOUT")))
+                response.raise_for_status()
+                image = open_image(BytesIO(response.content))
+                temp_file_path = save_image_to_temp_file(image)
+                classifications = client.image_classification(temp_file_path, model=getenv("IMAGE_CLASSIFICATION_MODEL"))
+            except Exception as e:
+                raise gr.Error(f"Failed to fetch image from URL: {str(e)}")
+        else:
+            raise gr.Error("Please either provide an image URL or upload an image.")
+        df = pd.DataFrame({
+                              "Label": classification.label,
+                              "Probability": f"{classification.score:.2%}"
+                          }
+                          for classification
+                          in classifications)
+        return image, df
+    finally:
+        # Clean up temporary file.
+        if temp_file_path and path.exists(temp_file_path):
+            try:
+                unlink(temp_file_path)
+            except Exception:
+                pass # Ignore clean-up errors.

text_to_image.py

@@ -0,0 +1,7 @@
+from os import getenv
+from PIL.Image import Image
+from huggingface_hub import InferenceClient
+
+
+def text_to_image(client: InferenceClient, prompt: str) -> Image:
+    return client.text_to_image(prompt, model=getenv("TEXT_TO_IMAGE_MODEL"))

utils.py

@@ -0,0 +1,12 @@
+from PIL.Image import Image
+from tempfile import NamedTemporaryFile
+
+
+def save_image_to_temp_file(image: Image) -> str:
+    image_format = image.format if image.format else 'PNG'
+    format_extension = image_format.lower() if image_format else 'png'
+    temp_file = NamedTemporaryFile(delete=False, suffix=f".{format_extension}")
+    temp_path = temp_file.name
+    temp_file.close()
+    image.save(temp_path, format=image_format)
+    return temp_path