txt2im-models / app.py
dn6's picture
dn6 HF staff
change name of comet logging module
8221e77
import uuid
import gradio as gr
import pandas as pd
from PIL import Image
from transformers import CLIPModel, CLIPProcessor
from comet import get_experiment, get_experiment_status, start_experiment
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 remember the good ones!
Set your Comet credentials in the Comet Settings tab and create an Experiment for logging data. If you don't have credentials yet,
you can [sign up for Comet here](https://www.comet.ml/signup)
If you want to continue logging to the same Experiment over multiple sessions, simply provide the experiment name.
Set a path to a Space using that uses a Diffusion model and submit your prompt in the Diffusion Model tab
** Note: ** This Space will still run even if you don't set credentials
"""
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 image, experiment_state
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("Start Experiment", variant="primary")
status = gr.Button("Experiment Status")
status_output = gr.Textbox(label="Status")
experiment_state = gr.Variable(label="Experiment State")
start.click(
start_experiment,
inputs=[
comet_api_key,
comet_workspace,
comet_project_name,
comet_experiment_name,
experiment_state,
],
outputs=[experiment_state, status_output],
)
status.click(
get_experiment_status,
inputs=[experiment_state],
outputs=[experiment_state, status_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",
placeholder="Enter a path to a Space",
)
prompt = gr.Textbox(
label="Prompt",
value="an oil painting of a corgi",
placeholder="Enter your text prompt here",
)
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()
start_interface()