app.py

"""Driver File."""

import gradio as gr
import core
import parameters

############################################ LAYOUT ############################################
with gr.Blocks() as demo:
    # ... [previous state variables remain the same] ...
    selected_image = gr.State(None)
    current_episode = gr.State(-1)
    current_scene = gr.State(-1)
    current_frame = gr.State(-1)
    episode_data = gr.State({})
    character_data = gr.State({})
    current_frame_data = gr.State(None)
    frame_hash_map = gr.State({})
    image_data_b64 = gr.State([])
    choice = gr.State([])

    with gr.Row():
        with gr.Column():
            comic_id = gr.Textbox(label="Enter Comic ID:", placeholder="Enter Comic ID")
        load_metadata = gr.Button("Load Metadata")

    # Display information about current Image
    with gr.Row():
        episode_dropdown = gr.Dropdown(
            choices=[], label="Current Episode", interactive=True
        )
        frame_dropdown = gr.Dropdown(
            choices=[], label="Current Frame", interactive=True
        )
        with gr.Column():
            load_images = gr.Button("Load Images")
            developer = gr.Checkbox(
                value=False, label="Enable Developer Mode"
            )

    images = gr.Gallery(
        label="Select an Image", elem_id="image_select", columns=4, height=320, object_fit="contain"
    )

    with gr.Row():
        image_description = gr.Textbox(label="Description", interactive=False)
        narration = gr.Textbox(label="narration", interactive=False)

    with gr.Row():
        character = gr.Textbox(label="Character", interactive=False)
        dialouge = gr.Textbox(label="Dialouge", interactive=False)
        location = gr.Textbox(label="Location", interactive=False)

    # buttons to interact with the data
    with gr.Row():
        prev_button = gr.Button("Prev Image")
        save_button = gr.Button("Save Image")
        next_button = gr.Button("Next Image")

    with gr.Column(visible=False) as developer_options:
        with gr.Column():
            setting = gr.Textbox(label="Frame Setting", interactive=True)
            with gr.Row():
                with gr.Column():
                    gr.Markdown("## Composition #1")
                    prompt_1 = gr.TextArea(label="Image Prompt", interactive=True)
                    seed_1 = gr.Textbox(label="Generation Seed")
                with gr.Column():
                    gr.Markdown("## Composition #2")
                    prompt_2 = gr.TextArea(label="Image Prompt", interactive=True)
                    seed_2 = gr.Textbox(label="Generation Seed")
            with gr.Row():
                with gr.Column():
                    gr.Markdown("## Composition #3")
                    prompt_3 = gr.TextArea(label="Image Prompt", interactive=True)
                    seed_3 = gr.Textbox(label="Generation Seed")
                with gr.Column():
                    gr.Markdown("## Composition #4")
                    prompt_4 = gr.TextArea(label="Image Prompt", interactive=True)
                    seed_4 = gr.Textbox(label="Generation Seed")

            with gr.Row():
                regenerate_comps_btn = gr.Button(value="Regenerate Compositions")

        with gr.Column():
            negative_prompt = gr.TextArea(
                value="",
                label="Negative Prompt",
                interactive=True,
            )
            related_chars = gr.CheckboxGroup(
                choices=[], value=[], label="Related Characters", interactive=True
            )

            with gr.Row():
                height = gr.Textbox(
                    value=parameters.IMG_HEIGHT, label="Image Height", interactive=True
                )
                width = gr.Textbox(
                    value=parameters.IMG_WIDTH, label="Image Width", interactive=True
                )
                visual_style = gr.Dropdown(
                    choices=parameters.VISUAL_CHOICES,
                    label="Current Visual Style",
                    interactive=True,
                )

        with gr.Row():
            regenerate_btn = gr.Button("Regenerate")
            save_btn = gr.Button("Save")

    ############################################ EVENTS ############################################
    developer.change(
        core.toggle_developer_options,
        inputs=[developer],
        outputs=[developer_options],
    ).then(
        lambda is_developer: (
            gr.update(interactive=is_developer),
            gr.update(interactive=is_developer),
            gr.update(interactive=is_developer),
            gr.update(interactive=is_developer),
            gr.update(interactive=is_developer),
        ),
        inputs=[developer],
        outputs=[image_description, narration, character, dialouge, location],
    )

    load_metadata.click(
        core.load_metadata_fn,
        inputs=[comic_id],
        outputs=[episode_dropdown, current_episode, character_data],
    ).then(  # This ensures `load_dropdown_fn` runs after `load_metadata_fn`
        core.episode_dropdown_effect,
        inputs=[comic_id, episode_dropdown],
        outputs=[
            frame_dropdown,
            current_episode,
            current_frame,
            episode_data,
            frame_hash_map,
        ],
    )

    episode_dropdown.input(
        core.episode_dropdown_effect,
        inputs=[comic_id, episode_dropdown],
        outputs=[
            frame_dropdown,
            current_episode,
            current_frame,
            episode_data,
            frame_hash_map,
        ],
    )

    frame_dropdown.input(
        lambda frame: frame,
        inputs=[frame_dropdown],
        outputs=[current_frame],
    )

    load_images.click(
        core.load_data,
        inputs=[
            episode_data,
            current_frame,
            frame_hash_map,
        ],
        outputs=[
            images,
            image_description,
            narration,
            character,
            dialouge,
            location,
            setting,
            prompt_1,
            seed_1,
            prompt_2,
            seed_2,
            prompt_3,
            seed_3,
            prompt_4,
            seed_4,
        ],
    ).then(
        core.update_characters,
        inputs=[character_data, current_frame, frame_hash_map, episode_data],
        outputs=[related_chars],
    )

    # When an image is clicked
    def get_select_index(evt: gr.SelectData, images):
        return images[evt.index]
    images.select(get_select_index, images, selected_image)

    next_button.click(
        core.load_data_next,
        inputs=[
            comic_id,
            current_episode,
            current_frame,
            frame_hash_map,
            episode_data,
        ],
        outputs=[
            episode_dropdown,
            frame_dropdown,
            current_episode,
            current_frame,
            episode_data,
            frame_hash_map,
        ],
    ).then(
        core.load_data,
        inputs=[
            episode_data,
            current_frame,
            frame_hash_map,
        ],
        outputs=[
            images,
            image_description,
            narration,
            character,
            dialouge,
            location,
            setting,
            prompt_1,
            seed_1,
            prompt_2,
            seed_2,
            prompt_3,
            seed_3,
            prompt_4,
            seed_4,
        ],
    ).then(
        core.update_characters,
        inputs=[character_data, current_frame, frame_hash_map, episode_data],
        outputs=[related_chars],
    ).then(
        lambda a, b: (a, b),
        inputs=[current_episode, current_frame],
        outputs=[current_episode, current_frame],
    )

    prev_button.click(
        core.load_data_prev,
        inputs=[
            comic_id,
            current_episode,
            current_frame,
            frame_hash_map,
            episode_data,
        ],
        outputs=[
            episode_dropdown,
            frame_dropdown,
            current_episode,
            current_frame,
            episode_data,
            frame_hash_map,
        ],
    ).then(
        core.load_data,
        inputs=[
            episode_data,
            current_frame,
            frame_hash_map,
        ],
        outputs=[
            images,
            image_description,
            narration,
            character,
            dialouge,
            location,
            setting,
            prompt_1,
            seed_1,
            prompt_2,
            seed_2,
            prompt_3,
            seed_3,
            prompt_4,
            seed_4,
        ],
    ).then(
        core.update_characters,
        inputs=[character_data, current_frame, frame_hash_map, episode_data],
        outputs=[related_chars],
    ).then(
        lambda a, b: (a, b),
        inputs=[current_episode, current_frame],
        outputs=[current_episode, current_frame],
    )

    save_button.click(
        core.save_image,
        inputs=[
            selected_image,
            comic_id,
            current_episode,
            current_frame,
        ],
        outputs=[],
    )

    regenerate_comps_btn.click(
        core.regenerate_compositions,
        inputs=[
            image_description,
            narration,
            character,
            dialouge,
            location,
            setting,
            related_chars,
            current_episode,
            current_frame,
            episode_data,
            frame_hash_map,
            character_data,
        ],
        outputs=[
            prompt_1,
            prompt_2,
            prompt_3,
            prompt_4,
        ],
    )

    regenerate_btn.click(
        core.regenerate_images,
        inputs=[
            current_episode,
            current_frame,
            visual_style,
            height,
            width,
            character_data,
            related_chars,
            prompt_1,
            seed_1,
            prompt_2,
            seed_2,
            prompt_3,
            seed_3,
            prompt_4,
            seed_4
        ],
        outputs=[images],
    )

    save_btn.click(
        core.save_comic_data,
        inputs=[
            current_episode,
            current_frame,
            episode_data,
            comic_id,
            image_description,
            narration,
            character,
            dialouge,
            location,
            setting,
            prompt_1,
            prompt_2,
            prompt_3,
            prompt_4,
            frame_hash_map,
            related_chars,
            character_data,
            images,
        ]
    )

demo.launch(
    # auth=("admin", "Qrt@12*34#immersfy"), share=True, ssr_mode=False, debug=True
    share=False,
)
# demo.launch(share=True, ssr_mode=False, debug=True)