Added character checkbox functionality

app.py

@@ -16,6 +16,7 @@ with gr.Blocks() as demo:
     current_frame_data = gr.State(None)
     details = gr.State({})
     image_data_b64 = gr.State([])
+    choice = gr.State([])
 
     with gr.Row():
         with gr.Column():
@@ -86,7 +87,11 @@ with gr.Blocks() as demo:
                 value="",
                 label="Negative Prompt",
             )
-            chars = gr.Textbox(value="[]", label="Related Characters")
+            # chars = gr.Textbox(value="[]", label="Related Characters")
+            chars = gr.CheckboxGroup(
+                choices=[], value=[], label="Related Characters", interactive=True
+            )
+
             with gr.Row():
                 height = gr.Textbox(value="1024", label="Image Height")
                 width = gr.Textbox(value="1024", label="Image Width")
@@ -137,7 +142,13 @@ with gr.Blocks() as demo:
 
     load_images.click(
         core.load_from_dropdown,
-        inputs=[episodes_data, current_episode, current_frame, developer],
+        inputs=[
+            episodes_data,
+            current_episode,
+            current_frame,
+            developer,
+            character_data,
+        ],
         outputs=[
             episode_dropdown,
             frame_dropdown,
@@ -172,7 +183,13 @@ with gr.Blocks() as demo:
 
     next_button.click(
         core.load_data_next,
-        inputs=[episodes_data, current_episode, current_frame, developer],
+        inputs=[
+            episodes_data,
+            current_episode,
+            current_frame,
+            developer,
+            character_data,
+        ],
         outputs=[
             episode_dropdown,
             frame_dropdown,
@@ -201,7 +218,13 @@ with gr.Blocks() as demo:
 
     prev_button.click(
         core.load_data_prev,
-        inputs=[episodes_data, current_episode, current_frame, developer],
+        inputs=[
+            episodes_data,
+            current_episode,
+            current_frame,
+            developer,
+            character_data,
+        ],
         outputs=[
             episode_dropdown,
             frame_dropdown,
@@ -312,20 +335,17 @@ with gr.Blocks() as demo:
             visual_style,
             height,
             width,
-            details
+            details,
+            chars,
         ],
         outputs=[images, image_data_b64],
     )
 
     save_btn.click(
         core.save_images,
-        inputs=[
-            image_data_b64,
-            current_episode,
-            current_frame,
-            comic_id,
-            details
-        ],
+        inputs=[image_data_b64, current_episode, current_frame, comic_id, details],
     )
-demo.launch(auth=("admin", "Qrt@12*34#immersfy"), share=True, ssr_mode=False,  debug=True)
+demo.launch(
+    auth=("admin", "Qrt@12*34#immersfy"), share=True, ssr_mode=False, debug=True
+)
 # demo.launch(share=True, ssr_mode=False, debug=True)

core.py

@@ -39,6 +39,12 @@ class ComicFrame:
     compositions: List[Composition] = dataclasses.field(default_factory=list)
 
 
+def update_characters(character_data, curr_frame):
+    return gr.CheckboxGroup(
+        choices=list(character_data.keys()), value=curr_frame.all_characters
+    )
+
+
 def list_current_dir(bucket_name: str, folder_path: str = "") -> list:
     try:
         response = aws_utils.S3_CLIENT.list_objects_v2(
@@ -54,7 +60,11 @@ def list_current_dir(bucket_name: str, folder_path: str = "") -> list:
 
 
 def load_data_inner(
-    episodes_data: list, current_episode: int, current_frame: int, is_developer: bool
+    episodes_data: list,
+    current_episode: int,
+    current_frame: int,
+    is_developer: bool,
+    character_data,
 ):
     try:
         images = []
@@ -81,7 +91,7 @@ def load_data_inner(
             curr_frame.compositions[2].seed,
             curr_frame.compositions[3].prompt,
             curr_frame.compositions[3].seed,
-            curr_frame.all_characters,
+            update_characters(character_data, curr_frame),
         )
     except Exception as e:
         return (
@@ -179,7 +189,11 @@ def load_metadata_fn(comic_id: str):
 
 
 def load_data_next(
-    episodes_data: list, current_episode: int, current_frame: int, is_developer: bool
+    episodes_data: list,
+    current_episode: int,
+    current_frame: int,
+    is_developer: bool,
+    character_data,
 ):
     if current_frame + 1 < len(episodes_data[current_episode]):
         current_frame += 1
@@ -191,12 +205,18 @@ def load_data_next(
     return (
         gr.update(value=current_episode),
         gr.update(value=current_frame),
-        *load_data_inner(episodes_data, current_episode, current_frame, is_developer),
+        *load_data_inner(
+            episodes_data, current_episode, current_frame, is_developer, character_data
+        ),
     )
 
 
 def load_data_prev(
-    episodes_data: list, current_episode: int, current_frame: int, is_developer: bool
+    episodes_data: list,
+    current_episode: int,
+    current_frame: int,
+    is_developer: bool,
+    character_data,
 ):
     if current_frame - 1 >= 0:
         current_frame -= 1
@@ -208,17 +228,30 @@ def load_data_prev(
     return (
         gr.update(value=current_episode),
         gr.update(value=current_frame),
-        *load_data_inner(episodes_data, current_episode, current_frame, is_developer),
+        *load_data_inner(
+            episodes_data, current_episode, current_frame, is_developer, character_data
+        ),
     )
 
 
 def load_from_dropdown(
-    episodes_data: dict, selected_episode: int, selected_frame: int, is_developer: bool
+    episodes_data: dict,
+    selected_episode: int,
+    selected_frame: int,
+    is_developer: bool,
+    character_data,
 ):
     return (
         gr.update(value=selected_episode),
         gr.update(value=selected_frame),
-        *load_data_inner(episodes_data, selected_episode, selected_frame, is_developer),
+        *load_data_inner(
+            episodes_data,
+            selected_episode,
+            selected_frame,
+            is_developer,
+            character_data,
+        ),
+        # update_characters(character_data)
     )
 
 
@@ -375,9 +408,11 @@ def regenerate_data(
     height,
     width,
     details,
+    chars,
 ):
     images = []
     image_data_b64 = []
+    print(chars)
     try:
         current_scene, frame_num_in_scene = get_scene_number(
             details, current_episode, current_frame
@@ -386,7 +421,7 @@ def regenerate_data(
             f"Regenerating data for episode {current_episode}, Scene: {current_scene}, frame {frame_num_in_scene}"
         )
         frame = episodes_data[current_episode][current_frame]
-        related_chars = [character_data[ch] for ch in frame.all_characters]
+        related_chars = [character_data[ch] for ch in chars]
 
         for i, composition in enumerate(frame.compositions):
             try:
@@ -525,5 +560,6 @@ def save_images(image_data_b64, current_episode, current_frame, comic_id, detail
                 print(f"Error saving image {i} to S3: {e}")
                 continue
         print("All images saved successfully!")
+        gr.Info("All images saved successfully!")
     except Exception as e:
         print(f"Error in save_images: {e}")