experiments/arena/pages/arena.py

# Copyright 2024 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from dataclasses import field
import random
import logging
import time
from concurrent.futures import ThreadPoolExecutor, as_completed

import mesop as me

from common.metadata import update_elo_ratings
from config.default import Default
from prompts.utils import PromptManager
from state.state import AppState
from components.header import header

from models.set_up import ModelSetup, load_default_models

from models.gemini_model import (
    generate_content,
    generate_images,
)
from models.generate import images_from_flux, images_from_imagen, images_from_stable_diffusion, study_fetch


# Initialize configuration
client, model_id = ModelSetup.init()
MODEL_ID = model_id
config = Default()
prompt_manager = PromptManager()
logging.basicConfig(level=logging.DEBUG)


IMAGEN_MODELS = [config.MODEL_IMAGEN2, config.MODEL_IMAGEN3_FAST, config.MODEL_IMAGEN3, config.MODEL_IMAGEN32,]
GEMINI_MODELS = [config.MODEL_GEMINI2]


@me.stateclass
class PageState:
    """Local Page State"""

    temp_name: str = ""
    is_loading: bool = False

    # pylint: disable=invalid-field-call
    arena_prompt: str = ""
    image_negative_prompt_input: str = ""
    image_aspect_ratio: str = "1:1"
    arena_textarea_key: int = 0
    arena_model1: str = ""
    arena_model2: str = ""
    arena_output: list[str] = field(default_factory=lambda: [])
    chosen_model: str = ""
    study: str = "live"
    study_models: list[str] = field(default_factory=list)
    # pylint: disable=invalid-field-call


def arena_images(input: str, study: str):
    """Create images for arena comparison"""
    state = me.state(PageState)
    if input == "":  # handle condition where someone hits "random" but doesn't modify
        if state.arena_prompt != "":
            input = state.arena_prompt
    state.arena_output.clear()

    logging.info("BATTLE: %s vs. %s", state.arena_model1, state.arena_model2)

    prompt = input
    logging.info("prompt: %s", prompt)
    if state.image_negative_prompt_input:
        logging.info("negative prompt: %s", state.image_negative_prompt_input)
    
    with ThreadPoolExecutor() as executor:  # Create a thread pool
        futures = []
        if study == "live":
            # model 1
            if state.arena_model1 in IMAGEN_MODELS:
                logging.info("model 1: %s", state.arena_model1)
                futures.append(
                    executor.submit(
                        images_from_imagen,
                        state.arena_model1,
                        prompt,
                        state.image_aspect_ratio,
                    )
                )
            elif state.arena_model1.startswith(config.MODEL_GEMINI2):
                logging.info("model 1: %s", state.arena_model1)
                futures.append(
                    executor.submit(
                        generate_images,
                        prompt,
                    )
                )
            elif state.arena_model1.startswith(config.MODEL_FLUX1):
                if config.MODEL_FLUX1_ENDPOINT_ID:
                    logging.info("model 1: %s", state.arena_model1)
                    futures.append(
                        executor.submit(
                            images_from_flux,
                            state.arena_model1,
                            prompt,
                            state.image_aspect_ratio,
                        )
                    )
                else:
                    logging.error("no endpoint defined for %s", state.arena_model1)
            elif state.arena_model1.startswith(config.MODEL_STABLE_DIFFUSION):
                if config.MODEL_STABLE_DIFFUSION_ENDPOINT_ID:
                    logging.info("model 1: %s", state.arena_model1)
                    futures.append(
                        executor.submit(
                            images_from_stable_diffusion,
                            state.arena_model1,
                            prompt,
                            state.image_aspect_ratio,
                        )
                    )
                else:
                    logging.error("no endpoint defined for %s", state.arena_model1)

            # model 2
            if state.arena_model2 in IMAGEN_MODELS:
                logging.info("model 2: %s", state.arena_model2)
                futures.append(
                    executor.submit(
                        images_from_imagen,
                        state.arena_model2,
                        prompt,
                        state.image_aspect_ratio,
                    )
                )
            elif state.arena_model2.startswith(config.MODEL_GEMINI2):
                logging.info("model 2: %s", state.arena_model2)
                futures.append(
                    executor.submit(
                        generate_images,
                        prompt,
                    )
                )
            elif state.arena_model2.startswith(config.MODEL_FLUX1):
                if config.MODEL_FLUX1_ENDPOINT_ID:
                    logging.info("model 2: %s", state.arena_model2)
                    futures.append(
                        executor.submit(
                            images_from_flux,
                            state.arena_model2,
                            prompt,
                            state.image_aspect_ratio,
                        )
                    )
                else:
                    logging.error("no endpoint defined for %s", state.arena_model2)
            elif state.arena_model2.startswith(config.MODEL_STABLE_DIFFUSION):
                if config.MODEL_STABLE_DIFFUSION_ENDPOINT_ID:
                    logging.info("model 2: %s", state.arena_model2)
                    futures.append(
                        executor.submit(
                            images_from_stable_diffusion,
                            state.arena_model2,
                            prompt,
                            state.image_aspect_ratio,
                        )
                    )
                else:
                    logging.error("no endpoint defined for %s", state.arena_model2)
        # Fetch images from study
        else:
            futures.extend([
                executor.submit(
                    study_fetch,
                    state.arena_model1,
                    prompt
                ),
                executor.submit(
                    study_fetch,
                    state.arena_model2,
                    prompt
                )
            ])
        
        for future in as_completed(futures):  # Wait for tasks to complete
            try:
                result = future.result()  # Get the result of each task
                state.arena_output.extend(
                    result
                )  # Assuming images_from_imagen returns a list
            except Exception as e:
                logging.error(f"Error during image generation: {e}")

def on_click_reload_arena(e: me.ClickEvent):  # pylint: disable=unused-argument
    """Reload arena handler"""
    state = me.state(PageState)
    if state.study == "live":
        state.study_models = load_default_models()

    state.arena_prompt = prompt_manager.random_prompt()

    state.arena_output.clear()

    state.is_loading = True
    yield
    print(f"Use {state.study_models}")

    # get random images
    state.arena_model1, state.arena_model2 = random.sample(state.study_models, 2)
    logging.info("%s vs. %s", state.arena_model1, state.arena_model2)
    arena_images(state.arena_prompt, state.study)

    state.is_loading = False
    yield


def on_click_arena_vote(e: me.ClickEvent):
    """Arena vote handler"""
    state = me.state(PageState)
    model_name = getattr(state, e.key)
    logging.info("user preferred %s: %s", e.key, model_name)
    state.chosen_model = model_name
    yield
    # update the elo ratings
    update_elo_ratings(state.arena_model1, state.arena_model2, model_name, state.arena_output, state.arena_prompt, state.study)
    yield
    time.sleep(int(Default.SHOW_RESULTS_PAUSE_TIME))
    yield
    # clear the output and reload
    state.arena_output.clear()
    state.chosen_model = ""
    state.arena_prompt = prompt_manager.random_prompt()
    state.arena_model1, state.arena_model2 = random.sample(state.study_models, 2)
    yield
    arena_images(state.arena_prompt, state.study)
    yield


WELCOME_PROMPT = """
Welcome the user to the battle of the generative media images, and encourage participation by asserting their voting on the images presented. 
This should be one or two sentences.
"""


def reload_welcome(e: me.ClickEvent):  # pylint: disable=unused-argument
    """Handle regeneration of welcome message event"""
    app_state = me.state(AppState)
    app_state.welcome_message = generate_welcome()
    yield


def generate_welcome() -> str:
    """Generate a nice welcome message with Gemini 2.0"""
    return generate_content(WELCOME_PROMPT)


def arena_page_content(app_state: me.state):
    """Arena Mesop Page"""

    page_state = me.state(PageState)
    prompt_manager.prompts_location = app_state.study_prompts_location
    page_state.study = app_state.study
    if page_state.study == "live":
        app_state.study_models = load_default_models()
    page_state.study_models = app_state.study_models
    print(f"======> Starting Page state study models: {page_state.study_models}")

    # TODO this is an initialization function that should be extracted
    if not app_state.welcome_message:
        app_state.welcome_message = generate_welcome()
    if not page_state.arena_prompt:
        page_state.arena_prompt = prompt_manager.random_prompt()
        page_state.arena_model1, page_state.arena_model2 = random.sample(app_state.study_models, 2)
        arena_images(page_state.arena_prompt, app_state.study)

    with me.box(
        style=me.Style(
            display="flex",
            flex_direction="column",
            height="100%",
        ),
    ):
        with me.box(
            style=me.Style(
                background=me.theme_var("background"),
                height="100%",
                overflow_y="scroll",
                margin=me.Margin(bottom=20),
            )
        ):
            with me.box(
                style=me.Style(
                    background=me.theme_var("background"),
                    padding=me.Padding(top=24, left=24, right=24, bottom=24),
                    display="flex",
                    flex_direction="column",
                )
            ):
                header("Arena" + (f" [Active Study: {app_state.study}]" if app_state.study != "live" else ""), "stadium")

                # welcome message
                with me.box(
                    style=me.Style(
                        flex_grow=1,
                        display="flex",
                        align_items="center",
                        justify_content="center",
                    ),
                    on_click=reload_welcome,
                ):
                    me.text(
                        app_state.welcome_message,
                        style=me.Style(
                            width="80vw",
                            font_size="12pt",
                            font_style="italic",
                            color="gray",
                        ),
                    )

                me.box(style=me.Style(height="16px"))

                with me.box(
                    style=me.Style(
                        margin=me.Margin(left="auto", right="auto"),
                        width="min(1024px, 100%)",
                        gap="24px",
                        flex_grow=1,
                        display="flex",
                        flex_wrap="wrap",
                        flex_direction="column",
                        align_items="center",
                    )
                ):
                    # Prompt
                    with me.box(
                        style=me.Style(
                            display="flex",
                            flex_direction="column",
                            align_items="center",
                            width="85%",
                        )
                    ):
                        me.text(
                            "Select the output you prefer for the given prompt",
                            style=me.Style(font_weight=500, font_size="20px", text_transform="uppercase"),
                        )
                        me.box(style=me.Style(height=16))
                        me.text(page_state.arena_prompt, style=me.Style(font_size="20pt"))

                    # Image outputs
                    with me.box(style=_BOX_STYLE):
                        if page_state.is_loading:
                            with me.box(
                                style=me.Style(
                                    display="grid",
                                    justify_content="center",
                                    justify_items="center",
                                )
                            ):
                                me.progress_spinner()
                        if len(page_state.arena_output) != 0:
                            with me.box(
                                style=me.Style(
                                    display="grid",
                                    justify_content="center",
                                    justify_items="center",
                                )
                            ):
                                # Generated images row
                                with me.box(
                                    style=me.Style(
                                        flex_wrap="wrap", display="flex", gap="15px"
                                    )
                                ):
                                    for idx, img in enumerate(page_state.arena_output, start=1):
                                        print(f"===> idx: {idx}, img: {img}")
                                        model_name = f"arena_model{idx}"
                                        model_value = getattr(page_state, model_name)

                                        replace_url = "https://storage.mtls.cloud.google.com/"
                                        if Default.PUBLIC_BUCKET:
                                            replace_url = "https://storage.googleapis.com/"
                                        img_url = img.replace(
                                            "gs://",
                                            replace_url
                                        )
                                        with me.box(
                                            style=me.Style(align_items="center", justify_content="center", display="flex", flex_direction="column"),
                                        ):
                                            image_border_style = me.Style(
                                                width="450px",
                                                margin=me.Margin(top=10),
                                                border_radius="35px",
                                            )
                                            if page_state.chosen_model:
                                                if page_state.chosen_model == model_value:
                                                    # green border
                                                    image_border_style = me.Style(
                                                        width="450px",
                                                        margin=me.Margin(top=10),
                                                        border_radius="35px",
                                                        border=me.Border().all(me.BorderSide(color="green", style="inset", width="5px"))
                                                    )
                                                else:
                                                    # opaque
                                                    image_border_style = me.Style(
                                                        width="450px",
                                                        margin=me.Margin(top=10),
                                                        border_radius="35px",
                                                        opacity=0.5,
                                                    )
                                            me.image(
                                                src=f"{img_url}",
                                                style=image_border_style,
                                            )
                                            
                                            if page_state.chosen_model:
                                                text_style = me.Style()
                                                if page_state.chosen_model == model_value:
                                                    text_style = me.Style(font_weight="bold")
                                                me.text(model_value, style=text_style)
                                            else:
                                                me.box(style=me.Style(height=18))

                                me.box(style=me.Style(height=15))

                                if len(page_state.arena_output) != 2:
                                    disabled_choice = True
                                else:
                                    disabled_choice = False

                                with me.box(
                                    style=me.Style(
                                        flex_direction="row",
                                        display="flex",
                                        gap=50,
                                    )
                                ):
                                    # left choice button
                                    with me.content_button(
                                        type="flat",
                                        key="arena_model1",
                                        on_click=on_click_arena_vote,
                                        disabled=disabled_choice,
                                    ):
                                        with me.box(
                                            style=me.Style(
                                                display="flex", align_items="center"
                                            )
                                        ):
                                            me.icon("arrow_left")
                                            me.text("left")
                                    # skip button
                                    me.button(
                                        label="skip",
                                        type="stroked",
                                        on_click=on_click_reload_arena,
                                    )
                                    # right choice button
                                    with me.content_button(
                                        type="flat",
                                        key="arena_model2",
                                        on_click=on_click_arena_vote,
                                        disabled=disabled_choice,
                                    ):
                                        with me.box(
                                            style=me.Style(
                                                display="flex", align_items="center"
                                            )
                                        ):
                                            me.text("right")
                                            me.icon("arrow_right")
                        else:
                            # skip button
                            me.button(
                                label="skip",
                                type="stroked",
                                on_click=on_click_reload_arena,
                            )
                    # show user choice
                    if page_state.chosen_model:
                        me.text(f"You voted {page_state.chosen_model}")


_BOX_STYLE = me.Style(
    flex_basis="max(480px, calc(50% - 48px))",
    background=me.theme_var("background"),
    border_radius=12,
    box_shadow=("0 3px 1px -2px #0003, 0 2px 2px #00000024, 0 1px 5px #0000001f"),
    padding=me.Padding(top=16, left=16, right=16, bottom=16),
    display="flex",
    flex_direction="column",
    width="100%",
)