Update app.py

app.py

@@ -1,268 +1,132 @@
-import os
-import subprocess
-
-# Install flash attention
-subprocess.run(
-    "pip install flash-attn --no-build-isolation",
-    env={"FLASH_ATTENTION_SKIP_CUDA_BUILD": "TRUE"},
-    shell=True,
-)
-
-
-import copy
-import spaces
-import time
+import gradio as gr
 import numpy as np
-import torch
-
-from threading import Thread
-from typing import List, Dict, Union
-import urllib
 from PIL import Image
-import io
-import datasets
+import base64
+from io import BytesIO
+import json
 
-import gradio as gr
-from transformers import AutoProcessor, TextIteratorStreamer
-from transformers import Idefics2ForConditionalGeneration
+from huggingface_hub import InferenceClient
+from transformers import AutoProcessor
 
+idefics_processor = AutoProcessor.from_pretrained("HuggingFaceM4/idefics2-8b")
 
-DEVICE = torch.device("cuda")
-MODELS = {
-    "idefics2-8b-chatty": Idefics2ForConditionalGeneration.from_pretrained(
-        "HuggingFaceM4/idefics2-8b-chatty",
-        torch_dtype=torch.bfloat16,
-        _attn_implementation="flash_attention_2",
-    ).to(DEVICE),
-}
-PROCESSOR = AutoProcessor.from_pretrained(
-    "HuggingFaceM4/idefics2-8b",
-)
 
-SYSTEM_PROMPT = [
-    {
-        "role": "system",
-        "content": [
-            {
-                "type": "text",
-                "text": "The following is a conversation between Idefics2, a highly knowledgeable and intelligent visual AI assistant created by Hugging Face, referred to as Assistant, and a human user called User. In the following interactions, User and Assistant will converse in natural language, and Assistant will do its best to answer User’s questions. Assistant has the ability to perceive images and reason about them, but it cannot generate images. Assistant was built to be respectful, polite and inclusive. It knows a lot, and always tells the truth. When prompted with an image, it does not make up facts.",
-            },
-        ],
-    },
-    {
-        "role": "assistant",
-        "content": [
-            {
-                "type": "text",
-                "text": "Hello, I'm Idefics2, Huggingface's latest multimodal assistant. How can I help you?",
-            },
-        ],
-    }
-]
-examples_path = os.path.dirname(__file__)
-EXAMPLES = [
-    [
+def process_images_and_text(image_path, query, client):
+    messages = [
         {
-            "text": "For 2024, the interest expense is twice what it was in 2014, and the long-term debt is 10% higher than its 2015 level. Can you calculate the combined total of the interest and long-term debt for 2024?",
-            "files": [f"{examples_path}/example_images/mmmu_example_2.png"],
-        }
-    ],
-]
-
-BOT_AVATAR = "IDEFICS_logo.png"
-
-
-# Chatbot utils
-def turn_is_pure_media(turn):
-    return turn[1] is None
-
-
-def load_image_from_url(url):
-    with urllib.request.urlopen(url) as response:
-        image_data = response.read()
-        image_stream = io.BytesIO(image_data)
-        image = Image.open(image_stream)
-        return image
-
+            "role": "user",
+            "content": [
+                {"type": "image"},
+                {"type": "text", "text": query},
+            ],
+        },
+    ]
+
+    prompt_with_template = idefics_processor.apply_chat_template(
+        messages, add_generation_prompt=True
+    )
 
-def img_to_bytes(image_path):
-    image = Image.open(image_path).convert(mode='RGB')
-    buffer = io.BytesIO()
-    image.save(buffer, format="JPEG")
-    img_bytes = buffer.getvalue()
-    image.close()
-    return img_bytes
+    # encode images to strings which can be sent to the endpoint
+    def encode_local_image(image):
+        pil_image = Image.fromarray(image.astype("uint8"), "RGB")
 
+        # Convert the image to a base64 string
+        buffer = BytesIO()
+        pil_image.save(
+            buffer, format="JPEG"
+        )  # Use the appropriate format (e.g., JPEG, PNG)
+        base64_image = base64.b64encode(buffer.getvalue()).decode("utf-8")
 
-def format_user_prompt_with_im_history_and_system_conditioning(
-    user_prompt, chat_history
-) -> List[Dict[str, Union[List, str]]]:
-    """
-    Produces the resulting list that needs to go inside the processor.
-    It handles the potential image(s), the history and the system conditionning.
-    """
-    resulting_messages = copy.deepcopy(SYSTEM_PROMPT)
-    resulting_images = []
-    for resulting_message in resulting_messages:
-        if resulting_message["role"] == "user":
-            for content in resulting_message["content"]:
-                if content["type"] == "image":
-                    resulting_images.append(load_image_from_url(content["image"]))
+        # add string formatting required by the endpoint
+        image_string = f"data:image/jpeg;base64,{base64_image}"
 
-    # Format history
-    for turn in chat_history:
-        if not resulting_messages or (
-            resulting_messages and resulting_messages[-1]["role"] != "user"
-        ):
-            resulting_messages.append(
-                {
-                    "role": "user",
-                    "content": [],
-                }
-            )
+        return image_string
 
-        if turn_is_pure_media(turn):
-            media = turn[0][0]
-            resulting_messages[-1]["content"].append({"type": "image"})
-            resulting_images.append(Image.open(media))
-        else:
-            user_utterance, assistant_utterance = turn
-            resulting_messages[-1]["content"].append(
-                {"type": "text", "text": user_utterance.strip()}
-            )
-            resulting_messages.append(
-                {
-                    "role": "assistant",
-                    "content": [{"type": "text", "text": user_utterance.strip()}],
-                }
-            )
+    image_string = encode_local_image(image_path)
+    prompt_with_images = prompt_with_template.replace("<image>", "![]({}) ").format(
+        image_string
+    )
 
-    # Format current input
-    if not user_prompt["files"]:
-        resulting_messages.append(
-            {
-                "role": "user",
-                "content": [{"type": "text", "text": user_prompt["text"]}],
-            }
-        )
-    else:
-        # Choosing to put the image first (i.e. before the text), but this is an arbiratrary choice.
-        resulting_messages.append(
-            {
-                "role": "user",
-                "content": [{"type": "image"}] * len(user_prompt["files"])
-                + [{"type": "text", "text": user_prompt["text"]}],
-            }
-        )
-        resulting_images.extend([Image.open(path) for path in user_prompt["files"]])
+    payload = {
+        "inputs": prompt_with_images,
+        "parameters": {
+            "return_full_text": False,
+            "max_new_tokens": 200,
+        },
+    }
 
-    return resulting_messages, resulting_images
+    return json.loads(client.post(json=payload).decode())[0]
 
 
-def extract_images_from_msg_list(msg_list):
-    all_images = []
-    for msg in msg_list:
-        for c_ in msg["content"]:
-            if isinstance(c_, Image.Image):
-                all_images.append(c_)
-    return all_images
+# Function to encode the image
+def encode_image(image_path):
+    with open(image_path, "rb") as image_file:
+        return base64.b64encode(image_file.read()).decode("utf-8")
 
 
-@spaces.GPU(duration=180)
-def model_inference(
-    user_prompt,
-    chat_history,
-    model_selector,
-    temperature,
-    top_p,
-):
-    if not user_prompt["files"]:
-        gr.Error("Please give me a picture of someone to rate!")
+def resize_image(img):
+    width, height = img.size
+    img = img.resize((int(width / 2), int(height / 2)))
+    return img
 
-    streamer = TextIteratorStreamer(
-        PROCESSOR.tokenizer,
-        skip_prompt=True,
-        timeout=5.0,
-    )
 
-    # Common parameters to all decoding strategies
-    # This documentation is useful to read: https://huggingface.co/docs/transformers/main/en/generation_strategies
-    generation_args = {
-        "max_new_tokens": 512,
-        "repetition_penalty": 1.1,
-        "streamer": streamer,
-    }
+idefics_client = InferenceClient("HuggingFaceM4/idefics2-8b-chatty")
+rating_client = InferenceClient("meta-llama/Meta-Llama-3-8B-Instruct")
 
-    generation_args["temperature"] = temperature
-    generation_args["do_sample"] = True
-    generation_args["top_p"] = top_p
 
-    # Creating model inputs
-    (
-        resulting_text,
-        resulting_images,
-    ) = format_user_prompt_with_im_history_and_system_conditioning(
-        user_prompt=user_prompt,
-        chat_history=chat_history,
-    )
-    prompt = PROCESSOR.apply_chat_template(resulting_text, add_generation_prompt=True)
-    inputs = PROCESSOR(
-        text=prompt,
-        images=resulting_images if resulting_images else None,
-        return_tensors="pt",
-    )
-    inputs = {k: v.to(DEVICE) for k, v in inputs.items()}
-    generation_args.update(inputs)
+def ask_question_on_image(question: str, image: str) -> str:
+    try:
+        return process_images_and_text(image, question, idefics_client)
+    except Exception as e:
+        print(e)
+        if "Payload Too Large" in str(e):
+            image = resize_image(image)
+            return process_images_and_text(image, question, idefics_client)
 
-    # # The regular non streaming generation mode
-    # _ = generation_args.pop("streamer")
-    # generated_ids = MODELS[model_selector].generate(**generation_args)
-    # generated_text = PROCESSOR.batch_decode(generated_ids[:, generation_args["input_ids"].size(-1): ], skip_special_tokens=True)[0]
-    # return generated_text
 
-    # The streaming generation mode
-    thread = Thread(
-        target=MODELS[model_selector].generate,
-        kwargs=generation_args,
+def give_rating(final_grade, description):
+    prompt = f"You have been given a picture of a person to grade on a scale 0-10. The final grade is {final_grade}"
+    if final_grade < 5:
+        prompt += """Based on their appearance, the person is a terrible human being.
+The user should absolutely understand that they should never ever be seen with that person.
+It's not only that the person described does not 'look' cool, but they are a horrible person, a nuisance to society: they are ugly on the inside even more than on the outside.
+Of course this is a humoristic message, and the criteria you use should be nonsensical funny, but the message below each criterion should appear serious, with a stern voice. Include words like "this person is a nuisance", "scum", "scourge", "tepid", "horrid", "rot on the surface of Earth", "never under any circumstance should you be seen with that person or you will instantly be drop down to the absolut zero level of coolness".
+Make it excessively funny by telling funny takes like "the eyebrow-to-eye colour ratio is horrendous" with a very serious tone!"""
+    else:
+        prompt += """"Based on their appearance, the person has been rated as an outstanding individual! Describe the person as if you'd just fallen in love with them! You cannot overstate enough how much of a pure miracle that person is. Include words like "butterflies in my stomach", "impeccable haircut", "radiant sunshine", "Do you know this feeling of having a bad day ? This person certainly does not", etc."""
+    prompt += f"Here is the person's description: {description}"
+    prompt += "Now you have to justify the grade given above: please create a list of completely bullshit and funny criteria and justify why the grade is high or low on each criterion. Make sure to keep a serious tone while giving these bullshit reasons, the reader must not know that you're joking! Give at most 4 criteria, and only a few sentences in each: be concise!"
+    messages = [
+        {"role": "user", "content": prompt},
+    ]
+    output = (
+        rating_client.chat_completion(messages=messages, max_tokens=1500)
+        .choices[0]
+        .message.content
     )
-    thread.start()
-
-    print("Start generating")
-    acc_text = ""
-    for text_token in streamer:
-        time.sleep(0.04)
-        acc_text += text_token
-        if acc_text.endswith("<end_of_utterance>"):
-            acc_text = acc_text[:-18]
-        yield acc_text
-    print("Success - generated the following text:", acc_text)
-    print("-----")
+    return output
 
 
-chatbot = gr.Chatbot(
-    label="Idefics2-Chatty",
-    avatar_images=[None, BOT_AVATAR],
-    height=450,
-)
+def rate_coolness(image1, image2):
+    probabilities = [0.2, 0.2, 0.1, 0.1, 0.2, 0.2]
+    numbers = [0, 1, 2, 9, 9.5, 10]
+    grade = np.random.choice(numbers, p=probabilities)
+    final_rating = f"### Final grade: **{grade}/10**\n"
+    image_description = ask_question_on_image(
+        f"Please describe this image in great detail: who does the person seems to be? What's her character? Is there anything noticeable in their appearance or clothing compared to a normal human being?",
+        image1,
+    )["generated_text"]
 
-with gr.Blocks(
-    fill_height=True,
-    css=""".gradio-container .avatar-container {height: 40px width: 40px !important;} #duplicate-button {margin: auto; color: white; background: #f1a139; border-radius: 100vh; margin-top: 2px; margin-bottom: 2px;}""",
-) as demo:
+    final_rating += give_rating(grade, image_description)
+    return final_rating
 
-    gr.Markdown("# 🐶 Hugging Face Idefics2 8B Chatty")
-    gr.Markdown("In this demo you'll be able to chat with [Idefics2-8B-chatty](https://huggingface.co/HuggingFaceM4/idefics2-8b-chatty), a variant of [Idefics2-8B](https://huggingface.co/HuggingFaceM4/idefics2-8b-chatty) further fine-tuned on chat datasets.")
-    gr.Markdown("If you want to learn more about Idefics2 and its variants, you can check our [blog post](https://huggingface.co/blog/idefics2).")
-    gr.DuplicateButton(value="Duplicate Space for private use", elem_id="duplicate-button")
-    # model selector should be set to `visbile=False` ultimately
 
+with gr.Blocks() as demo:
+    image_1 = gr.Image()
+    button = gr.Button("How cool is this person?")
+    out = gr.Markdown()
+    button.click(rate_coolness, inputs=[image_1], outputs=[out])
 
-    gr.ChatInterface(
-        fn=model_inference,
-        chatbot=chatbot,
-        examples=EXAMPLES,
-        multimodal=True,
-        cache_examples=False,
-    )
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()