refactor to use gradio state

app.py

@@ -1,123 +1,187 @@
-import comet_ml
+import json
 import uuid
+
+import comet_ml
 import gradio as gr
+import pandas as pd
+from PIL import Image
+from transformers import CLIPModel, CLIPProcessor
+
+CLIP_MODEL_PATH = "openai/clip-vit-base-patch32"
 
+clip_model = CLIPModel.from_pretrained(CLIP_MODEL_PATH)
+clip_processor = CLIPProcessor.from_pretrained(CLIP_MODEL_PATH)
 
 DESCRIPTION = """Glad to see you here 😄.
 You can use this Space to log predictions to [Comet](https://www.comet.ml/site) from Spaces that use Text to Image Diffusion Models.
 
-Keep track of all your prompts and generated images so that you can easily find the good ones!
+Keep track of all your prompts and generated images so that you remember the good ones!
 
-Set your Comet credentials in the Comet Settings tab, and select the Space you want to log predictions
-from in the Diffusion Model tab
-"""
+Set your Comet credentials in the Comet Settings tab and create an Experiment for logging data.
+If you want to continue logging to the same Experiment over multiple sessions, add in the
 
+Then use the path to a Space to generate from in the Diffusion Model tab
+"""
 
-class MyProject:
-    def __init__(self):
-        self.experiment = None
 
-    def start_experiment(
-        self, comet_api_key: str, comet_workspace: str, comet_project_name: str
-    ):
-        if not comet_api_key:
-            return """
+def create_experiment(
+    comet_api_key,
+    comet_workspace,
+    comet_project_name,
+    comet_experiment_name,
+    experiment,
+):
+    if not comet_api_key:
+        experiment = None
+        return (
+            experiment,
+            """
             Please add your API key in order to log your predictions to a Comet Experiment.
             If you don't have a Comet account yet, you can sign up using the link below:
 
-            [Sign Up for Comet](https://www.comet.ml/signup)
-            """
+            https://www.comet.ml/signup
+            """,
+        )
 
-        else:
-            try:
-                self.experiment = comet_ml.Experiment(
-                    api_key=comet_api_key,
-                    workspace=comet_workspace,
-                    project_name=comet_project_name,
-                )
-                self.experiment.add_tags(["spaces"])
-
-                return f"Started {self.experiment.name}. Happy logging!😊"
-
-            except Exception as e:
-                return e
-
-    def end_experiment(self):
-        if self.experiment is not None:
-            self.experiment.end()
-            return f"Ended {self.experiment.name}"
-
-    def get_experiment_status(self):
-        return f"Running {self.experiment.name}"
-
-    def start_comet_interface(self):
-        demo = gr.Blocks()
-        with demo:
-            # credentials
-            comet_api_key = gr.Textbox(label="Comet API Key")
-            comet_workspace = gr.Textbox(label="Comet Workspace")
-            comet_project_name = gr.Textbox(label="Comet Project Name")
-
-            with gr.Row():
-                start_experiment = gr.Button("Start Experiment", variant="primary")
-                status = gr.Button("Experiment Status")
-                end_experiment = gr.Button("End Experiment", variant="secondary")
-
-            output = gr.Markdown(label="Status")
-
-            start_experiment.click(
-                self.start_experiment,
-                inputs=[
-                    comet_api_key,
-                    comet_workspace,
-                    comet_project_name,
-                ],
-                outputs=output,
-            )
-            status.click(self.get_experiment_status, inputs=None, outputs=None)
-            end_experiment.click(self.end_experiment, inputs=None, outputs=output)
-
-        return demo
-
-    def predict(
-        self,
-        model,
-        prompt,
-    ):
-        io = gr.Interface.load(model)
-        image = io(prompt)
-
-        if self.experiment is not None:
-            image_id = uuid.uuid4().hex
-            self.experiment.log_image(image, name=image_id)
-            self.experiment.log_text(
-                prompt, metadata={"image_id": image_id, "model": model}
-            )
-
-        return image
-
-    def load_interface(self):
-        model = gr.Textbox(label="Model", value="spaces/valhalla/glide-text2im")
-        prompt = gr.Textbox(label="Prompt")
-
-        outputs = gr.Image(label="Image")
-
-        interface = gr.Interface(
-            self.predict,
-            inputs=[model, prompt],
-            outputs=outputs,
-            examples=[["spaces/valhalla/glide-text2im", "An oil painting of a corgi"]],
-            description=DESCRIPTION,
+    try:
+        api_experiment = comet_ml.APIExperiment(
+            api_key=comet_api_key,
+            workspace=comet_workspace,
+            project_name=comet_project_name,
+            experiment_name=comet_experiment_name,
         )
+        experiment = {
+            "api_key": comet_api_key,
+            "workspace": comet_workspace,
+            "project_name": comet_project_name,
+            "previous_experiment": api_experiment.id,
+        }
+
+        return experiment, f"Started {api_experiment.name}. Happy logging!😊"
+
+    except Exception as e:
+        return None, e
+
+
+def get_experiment(kwargs) -> comet_ml.APIExperiment:
+    try:
+        return comet_ml.APIExperiment(**kwargs)
+    except Exception as e:
+        return None
+
+
+def get_experiment_status(experiment_state):
+    experiment = get_experiment(experiment_state)
+    if experiment is not None:
+        name = experiment.name
+        return experiment_state, f"Currently logging to: {name}"
+
+    return experiment_state, f"No Experiments found"
+
+
+def predict(
+    model,
+    prompt,
+    experiment_state,
+):
+    io = gr.Interface.load(model)
+    image = io(prompt)
+    pil_image = Image.open(image)
+
+    inputs = clip_processor(
+        text=[prompt],
+        images=pil_image,
+        return_tensors="pt",
+        padding=True,
+    )
+    outputs = clip_model(**inputs)
+    clip_score = outputs.logits_per_image.item() / 100.0
+
+    experiment = get_experiment(experiment_state)
+    if experiment is not None:
+        image_id = uuid.uuid4().hex
+        experiment.log_image(image, image_id)
+
+        asset = pd.DataFrame.from_records(
+            [
+                {
+                    "prompt": prompt,
+                    "model": model,
+                    "clip_model": CLIP_MODEL_PATH,
+                    "clip_score": round(clip_score, 3),
+                }
+            ]
+        )
+        experiment.log_table(f"{image_id}.json", asset, orient="records")
 
-        return interface
+    return image, experiment_state
 
-    def launch(self):
-        interface = gr.TabbedInterface(
-            [self.start_comet_interface(), self.load_interface()],
-            tab_names=["Comet Settings", "Diffusion Model"],
-        )
-        interface.launch()
+
+def start_interface():
+    demo = gr.Blocks()
+    with demo:
+        description = gr.Markdown(DESCRIPTION)
+        with gr.Tabs():
+            with gr.TabItem(label="Comet Settings"):
+                # credentials
+                comet_api_key = gr.Textbox(
+                    label="Comet API Key",
+                    placeholder="This is required if you'd like to create an Experiment",
+                )
+                comet_workspace = gr.Textbox(label="Comet Workspace")
+                comet_project_name = gr.Textbox(label="Comet Project Name")
+                comet_experiment_name = gr.Textbox(
+                    label="Comet Experiment Name",
+                    placeholder=(
+                        "Set this if you'd like"
+                        "to continue logging to an existing Experiment",
+                    ),
+                )
+
+                with gr.Row():
+                    start = gr.Button("Create Experiment", variant="primary")
+                    status = gr.Button("Experiment Status")
+
+                output = gr.Markdown(label="Status")
+                experiment_state = gr.Variable(label="Experiment State")
+
+                start.click(
+                    create_experiment,
+                    inputs=[
+                        comet_api_key,
+                        comet_workspace,
+                        comet_project_name,
+                        comet_experiment_name,
+                        experiment_state,
+                    ],
+                    outputs=[experiment_state, output],
+                )
+
+                status.click(
+                    get_experiment_status,
+                    inputs=[experiment_state],
+                    outputs=[experiment_state, output],
+                )
+
+            with gr.TabItem(label="Diffusion Model"):
+                diff_description = gr.Markdown(
+                    """The Model must be a path to any Space that accepts"
+                    only text as input and produces an image as an output
+                    """
+                )
+                model = gr.Textbox(label="Model", value="spaces/valhalla/glide-text2im")
+                prompt = gr.Textbox(label="Prompt")
+
+                outputs = gr.Image(label="Image")
+
+                submit = gr.Button("Submit", variant="primary")
+                submit.click(
+                    predict,
+                    inputs=[model, prompt, experiment_state],
+                    outputs=[outputs, experiment_state],
+                )
+
+    demo.launch()
 
 
-MyProject().launch()
+start_interface()