def image_generation_cols()

in genai-design-marketing-studio/streamlit_app.py [0:0]


def image_generation_cols(nb_image: bool=True, allow_kids: bool=False) -> None:
    """
    Image generation parameters columns
    :param nb_image: add a column for choosing how many images are generated if True
    :param allow_kids: add the allow_all option if True
    :return: None
    """

    nb_columns = 4 if nb_image else 3
    cols = st.columns(nb_columns)
    i = 0

    if nb_image:
        with cols[i]:
            number_of_images = st.number_input(
                "Number of images",
                min_value=1,
                max_value=4
            )
        i += 1
    else:
        number_of_images = 1

    with cols[i]:
        aspect_ratio = st.selectbox(
            "Aspect Ratio",
            options=["1:1", "4:3", "3:4", "9:16", "16:9"],
            index=0
        )
    i += 1

    with cols[i]:
        options = ["allow_adult", "dont_allow"]
        if allow_kids:
            options.prepend("allow_all")
        person_safety = st.selectbox(
            "Allow Person",
            options=options,
            index=0
        )
    i += 1

    with cols[i]:
        watermarking = st.checkbox("Watermark")

    return number_of_images, aspect_ratio, person_safety, watermarking