from transformers import pipeline
import gradio as gr


clip_models = [
  "zer0int/CLIP-GmP-ViT-L-14",
  "openai/clip-vit-large-patch14",
  "laion/CLIP-ViT-bigG-14-laion2B-39B-b160k",
]

clip_checkpoint = clip_models[0]
clip_detector = pipeline(model=clip_checkpoint, task="zero-shot-image-classification")

def postprocess(output):
  return {out["label"]: float(out["score"]) for out in output}


def infer(image, candidate_labels):
  candidate_labels = [label.lstrip(" ") for label in candidate_labels.split(",")]
  clip_out = clip_detector(image, candidate_labels=candidate_labels)
  return postprocess(clip_out)


def load_clip_model(modelname):
  global clip_detector
  try:
    clip_detector = pipeline(model=modelname, task="zero-shot-image-classification")
  except Exception as e:
    raise gr.Error(f"Model load error: {modelname} {e}")
  return modelname


with gr.Blocks() as demo:
  gr.Markdown("# Test CLIP")
  with gr.Row():
    with gr.Column():
        image_input = gr.Image(type="pil")
        text_input = gr.Textbox(label="Input a list of labels")
        model_input = gr.Dropdown(label="CLIP model", choices=clip_models, value=clip_models[0], allow_custom_value=True, interactive=True)
        run_button = gr.Button("Run", visible=True)

    with gr.Column():
      clip_output = gr.Label(label = "CLIP Output", num_top_classes=3)
      
  examples = [["./baklava.jpg", "baklava, souffle, tiramisu"], ["./cheetah.jpg", "cat, dog"], ["./cat.png", "cat, dog"]]
  gr.Examples(
        examples = examples, 
        inputs=[image_input, text_input],
        outputs=[clip_output],
        fn=infer,
        cache_examples=True
    )
  run_button.click(fn=infer,
                    inputs=[image_input, text_input],
                    outputs=[clip_output])
  model_input.change(load_clip_model, [model_input], [model_input])

demo.launch()