Major Speed update

.gitignore

@@ -1,2 +1,3 @@
-.env
-__pycache__
+*.env
+__pycache__/
+.gradio/

app.py

@@ -11,10 +11,10 @@ with gr.Blocks() as demo:
     current_episode = gr.State(-1)
     current_scene = gr.State(-1)
     current_frame = gr.State(-1)
-    episodes_data = gr.State({})
+    episode_data = gr.State({})
     character_data = gr.State({})
     current_frame_data = gr.State(None)
-    details = gr.State({})
+    frame_hash_map = gr.State({})
     image_data_b64 = gr.State([])
     choice = gr.State([])
 
@@ -34,11 +34,11 @@ with gr.Blocks() as demo:
         with gr.Column():
             load_images = gr.Button("Load Images")
             developer = gr.Checkbox(
-                value=False, label="Enable Developer Mode", visible=False
+                value=False, label="Enable Developer Mode"
             )
 
     images = gr.Gallery(
-        label="Select an Image", elem_id="image_select", columns=4, height=300
+        label="Select an Image", elem_id="image_select", columns=4, height=320, object_fit="contain"
     )
 
     with gr.Row():
@@ -58,43 +58,46 @@ with gr.Blocks() as demo:
 
     with gr.Column(visible=False) as developer_options:
         with gr.Column():
-            setting = gr.Textbox(label="Frame Setting")
+            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")
+                    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")
+                    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")
+                    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")
+                    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")
-                save_components_btn = gr.Button(value="Save Compositions")
 
         with gr.Column():
             negative_prompt = gr.TextArea(
                 value="",
                 label="Negative Prompt",
+                interactive=True,
             )
-            # chars = gr.Textbox(value="[]", label="Related Characters")
-            chars = gr.CheckboxGroup(
+            related_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")
+                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",
@@ -106,57 +109,65 @@ with gr.Blocks() as demo:
             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=[
-            episode_dropdown,
             frame_dropdown,
             current_episode,
             current_frame,
-            episodes_data,
-            character_data,
-            details,
-            developer,
+            episode_data,
+            frame_hash_map,
         ],
     )
 
     episode_dropdown.input(
-        core.load_dropdown_fn,
-        inputs=[episode_dropdown],
+        core.episode_dropdown_effect,
+        inputs=[comic_id, episode_dropdown],
         outputs=[
-            episode_dropdown,
             frame_dropdown,
             current_episode,
             current_frame,
+            episode_data,
+            frame_hash_map,
         ],
     )
 
     frame_dropdown.input(
-        core.load_dropdown_fn_v2,
+        lambda frame: frame,
         inputs=[frame_dropdown],
-        outputs=[
-            current_frame,
-        ],
+        outputs=[current_frame],
     )
 
     load_images.click(
-        core.load_from_dropdown,
+        core.load_data,
         inputs=[
-            episodes_data,
-            current_episode,
+            episode_data,
             current_frame,
-            developer,
-            character_data,
+            frame_hash_map,
         ],
         outputs=[
-            episode_dropdown,
-            frame_dropdown,
-            # Textual data returned from load_text_data()
             images,
-            episodes_data,
-            current_episode,
-            current_frame,
             image_description,
             narration,
             character,
@@ -171,33 +182,44 @@ with gr.Blocks() as demo:
             seed_3,
             prompt_4,
             seed_4,
-            chars,
         ],
+    ).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=[
-            episodes_data,
+            comic_id,
             current_episode,
             current_frame,
-            developer,
-            character_data,
+            frame_hash_map,
+            episode_data,
         ],
         outputs=[
             episode_dropdown,
             frame_dropdown,
-            # Textual data returned from load_text_data()
-            images,
-            episodes_data,
             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,
@@ -212,27 +234,43 @@ with gr.Blocks() as demo:
             seed_3,
             prompt_4,
             seed_4,
-            chars,
         ],
+    ).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=[
-            episodes_data,
+            comic_id,
             current_episode,
             current_frame,
-            developer,
-            character_data,
+            frame_hash_map,
+            episode_data,
         ],
         outputs=[
             episode_dropdown,
             frame_dropdown,
-            # Textual data returned from load_text_data()
-            images,
-            episodes_data,
             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,
@@ -247,8 +285,15 @@ with gr.Blocks() as demo:
             seed_3,
             prompt_4,
             seed_4,
-            chars,
         ],
+    ).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(
@@ -262,23 +307,8 @@ with gr.Blocks() as demo:
         outputs=[],
     )
 
-    developer.change(
-        core.toggle_developer_options,
-        inputs=[developer, prompt_1, prompt_2, prompt_3, prompt_4, setting],
-        outputs=[developer_options, prompt_1, prompt_2, prompt_3, prompt_4, setting],
-    ).then(
-        lambda: (
-            gr.update(interactive=True),
-            gr.update(interactive=True),
-            gr.update(interactive=True),
-            gr.update(interactive=True),
-            gr.update(interactive=True),
-        ),
-        outputs=[prompt_1, prompt_2, prompt_3, prompt_4, setting],
-    )
-
     regenerate_comps_btn.click(
-        core.regenerate_composition_data,
+        core.regenerate_compositions,
         inputs=[
             image_description,
             narration,
@@ -286,12 +316,31 @@ with gr.Blocks() as demo:
             dialouge,
             location,
             setting,
-            chars,
+            related_chars,
             current_episode,
             current_frame,
-            episodes_data,
+            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,
@@ -299,16 +348,17 @@ with gr.Blocks() as demo:
             prompt_3,
             seed_3,
             prompt_4,
-            seed_4,
+            seed_4
         ],
+        outputs=[images],
     )
 
-    save_components_btn.click(
-        core.save_image_compositions,
+    save_btn.click(
+        core.save_comic_data,
         inputs=[
             current_episode,
             current_frame,
-            details,
+            episode_data,
             comic_id,
             image_description,
             narration,
@@ -316,36 +366,19 @@ with gr.Blocks() as demo:
             dialouge,
             location,
             setting,
-            chars,
             prompt_1,
             prompt_2,
             prompt_3,
             prompt_4,
-        ],
-        outputs=[current_scene],
-    )
-    regenerate_btn.click(
-        core.regenerate_data,
-        inputs=[
-            comic_id,
-            current_episode,
-            current_frame,
-            episodes_data,
+            frame_hash_map,
+            related_chars,
             character_data,
-            visual_style,
-            height,
-            width,
-            details,
-            chars,
-        ],
-        outputs=[images, image_data_b64],
+            images,
+        ]
     )
 
-    save_btn.click(
-        core.save_images,
-        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
+    # auth=("admin", "Qrt@12*34#immersfy"), share=True, ssr_mode=False, debug=True
+    share=False,
 )
 # demo.launch(share=True, ssr_mode=False, debug=True)

aws_utils.py

@@ -84,7 +84,7 @@ def fetch_from_s3(source: Union[str, dict], region_name: str = "ap-south-1") ->
     Returns:
         bytes: The content of the file fetched from S3.
     """
-    print(f"Fetching file from S3. Source: {source}")
+    # print(f"Fetching file from S3. Source: {source}")
     s3_client = boto3.client("s3", region_name=region_name)
 
     # Parse the source depending on its type
@@ -104,7 +104,7 @@ def fetch_from_s3(source: Union[str, dict], region_name: str = "ap-south-1") ->
         print("Source must be a string URL or a dictionary.")
         raise ValueError("Source must be a string URL or a dictionary.")
 
-    print(f"Attempting to download from bucket: {bucket_name}, path: {file_path}")
+    # print(f"Attempting to download from bucket: {bucket_name}, path: {file_path}")
     try:
         response = s3_client.get_object(Bucket=bucket_name, Key=file_path)
         file_content = response["Body"].read()

core.py

@@ -1,7 +1,7 @@
 from typing import List
 from PIL import Image
+import json
 import gradio as gr
-import dataclasses
 import io
 import jinja2
 import base64
@@ -11,38 +11,20 @@ import script_gen
 import inout as iowrapper
 import openai_wrapper
 import json
-from dataclasses import asdict
 import base64
-import io
 
 AWS_BUCKET = parameters.AWS_BUCKET
 llm = openai_wrapper.GPT_4O_MINI
 
 
-@dataclasses.dataclass
-class Composition:
-    prompt: str
-    shot_type: str
-    seed: int
-    image: str
-
-
-@dataclasses.dataclass
-class ComicFrame:
-    description: str
-    narration: str
-    character_dilouge: str
-    character: str
-    location: str
-    setting: str
-    all_characters: list
-    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
-    )
+#### Functions ordered by their order of developement.
+def toggle_developer_options(is_developer: bool):
+    if is_developer:
+        # Return visibility updates for the developer options along with the values
+        return gr.update(visible=True)
+    else:
+        # Hide the developer options and return only the updated visibility
+        return gr.update(visible=False)
 
 
 def list_current_dir(bucket_name: str, folder_path: str = "") -> list:
@@ -59,227 +41,183 @@ def list_current_dir(bucket_name: str, folder_path: str = "") -> list:
         return []
 
 
-def load_data_inner(
-    episodes_data: list,
-    current_episode: int,
-    current_frame: int,
-    is_developer: bool,
-    character_data,
-):
+def load_metadata_fn(comic_id: str):
     try:
-        images = []
-        curr_frame = episodes_data[current_episode][current_frame]
-        for comp in curr_frame.compositions:
-            data = aws_utils.fetch_from_s3(comp.image)
-            images.append(Image.open(io.BytesIO(data)))
+        # Load character data
+        character_data = {}
+        character_path = f"s3://blix-demo-v0/{comic_id}/characters/characters.json"
+        char_data = json.loads(aws_utils.fetch_from_s3(character_path).decode("utf-8"))
+        character_data = {
+            name: char for name, char in char_data.items()
+        }
+
+        # Load episode data
+        episode_folders = list_current_dir(AWS_BUCKET, f"{comic_id}/episodes/")
+        episode_indices = []
+        for folder in episode_folders:
+            if "episode" in folder:
+                idx = int(folder.split("/")[2].split("-")[-1])
+                episode_indices.append(idx)
+        if not episode_indices:
+            return (gr.update(choices=[]), None, {})
+
+        # Return the values
+        min_episode = min(episode_indices)
         return (
-            images,
-            episodes_data,
-            current_episode,
-            current_frame,
-            gr.Textbox(value=curr_frame.description, interactive=is_developer),
-            gr.Textbox(value=curr_frame.narration, interactive=is_developer),
-            gr.Textbox(value=curr_frame.character, interactive=is_developer),
-            gr.Textbox(value=curr_frame.character_dilouge, interactive=is_developer),
-            gr.Textbox(value=curr_frame.location, interactive=is_developer),
-            curr_frame.setting,
-            curr_frame.compositions[0].prompt,
-            curr_frame.compositions[0].seed,
-            curr_frame.compositions[1].prompt,
-            curr_frame.compositions[1].seed,
-            curr_frame.compositions[2].prompt,
-            curr_frame.compositions[2].seed,
-            curr_frame.compositions[3].prompt,
-            curr_frame.compositions[3].seed,
-            update_characters(character_data, curr_frame),
+            gr.update(choices=episode_indices, value=min_episode),
+            min_episode,
+            character_data,
         )
     except Exception as e:
-        return (
-            [],
-            episodes_data,
-            current_episode,
-            current_frame,
-            gr.Textbox(),
-            gr.Textbox(),
-            gr.Textbox(),
-            gr.Textbox(),
-            gr.Textbox(),
-            "",
-            "",
-            "",
-            "",
-            "",
-            "",
-            "",
-            "",
-            "",
-            [],
+        gr.Warning(f"Error loading metadata: {e}")
+        return (gr.update(choices=[]), None, {})
+
+
+def load_episode_data(comic_id: str, episode_num: int):
+    try:
+        print(f"For episode: {episode_num}")
+        json_path = (
+            f"s3://{AWS_BUCKET}/{comic_id}/episodes/episode-{episode_num}/episode.json"
+        )
+        episode_data = json.loads(aws_utils.fetch_from_s3(json_path).decode("utf-8"))
+        frame_hash_map = {}
+        count = 1
+        for scene_idx, scene in enumerate(episode_data["scenes"]):
+            for frame_idx, _ in enumerate(scene["frames"]):
+                frame_hash_map[count] = {
+                    "scene": scene_idx,
+                    "frame": frame_idx,
+                }
+                count += 1
+        return (episode_data, frame_hash_map)
+    except Exception as e:
+        print(
+            f"Failed to load json dictionary for episode: {episode_num} at path: {json_path}"
         )
+        import traceback as tc
+        print(tc.format_exc())
+        return {}, {}
 
 
-def load_metadata_fn(comic_id: str):
+def episode_dropdown_effect(comic_id, selected_episode):
+    episode_data, frame_hash_map = load_episode_data(comic_id, selected_episode)
+    current_frame = min(list(frame_hash_map.keys()))
+    return (
+        gr.update(choices=list(frame_hash_map.keys()), value=current_frame),
+        selected_episode,
+        current_frame,
+        episode_data,
+        frame_hash_map,
+    )
+
+
+def load_data(episodes_data: dict, current_frame: int, frame_hash_map: dict):
     try:
-        episodes_data = {}
-        episode_idx = []
-        character_data = {}
-        details = {}
-        character_path = f"s3://blix-demo-v0/{comic_id}/characters/characters.json"
-        char_data = eval(aws_utils.fetch_from_s3(source=character_path).decode("utf-8"))
+        image_list = []
+        scene_num, frame_num = (
+            frame_hash_map[current_frame]["scene"],
+            frame_hash_map[current_frame]["frame"],
+        )
+        curr_frame = episodes_data["scenes"][scene_num]["frames"][frame_num]
 
-        for name, char in char_data.items():
-            character_data[name] = char["profile_image"]
+        for comp in curr_frame["compositions"]:
+            # Fetch image from S3
+            data = aws_utils.fetch_from_s3(comp["image"])
+            if data:
+                image = Image.open(io.BytesIO(data))
+                image_list.append(image)
+            else:
+                print(f"Failed to load image from: {comp['image']}")
 
-        for folder in list_current_dir(AWS_BUCKET, f"{comic_id}/episodes/"):
-            if "episode" in folder:
-                json_path = f"s3://{AWS_BUCKET}/{folder}episode.json"
-                idx = int(folder.split("/")[2].split("-")[-1])
-                episode_idx.append(idx)
-                data = eval(aws_utils.fetch_from_s3(source=json_path).decode("utf-8"))
-                comic_frames = []
-                details[idx] = {}
-                cumulative_frame_count = 0
-
-                for scene_num, scene in enumerate(data["scenes"]):
-                    scene_frame_count = len(scene["frames"])
-                    cumulative_frame_count += scene_frame_count
-                    details[idx][scene_num] = cumulative_frame_count
-
-                    for frame in scene["frames"]:
-                        comic_frames.append(
-                            ComicFrame(
-                                description=frame["description"],
-                                narration=frame["narration"],
-                                character=frame["audio_cue_character"],
-                                character_dilouge=frame["audio_cue_text"],
-                                compositions=[
-                                    Composition(**comp)
-                                    for comp in frame["compositions"]
-                                ],
-                                location=frame["location"],
-                                setting=frame["frame_setting"],
-                                all_characters=[
-                                    char["name"] for char in frame["characters"]
-                                ],
-                            )
-                        )
-                episodes_data[idx] = comic_frames
-
-        current_episode, current_frame = min(episode_idx), 0
         return (
-            gr.update(choices=episode_idx, value=episode_idx[0]),
-            gr.update(
-                choices=range(len(episodes_data[current_episode])), value=current_frame
-            ),
-            current_episode,
-            current_frame,
-            episodes_data,
-            character_data,
-            details,
-            gr.Checkbox(visible=True),
+            image_list,  # Return the image list to be displayed in the gallery
+            curr_frame["description"],
+            curr_frame["narration"],
+            curr_frame["audio_cue_character"],
+            curr_frame["audio_cue_text"],
+            curr_frame["location"],
+            curr_frame["frame_setting"],
+            curr_frame["compositions"][0]["prompt"],
+            curr_frame["compositions"][0]["seed"],
+            curr_frame["compositions"][1]["prompt"],
+            curr_frame["compositions"][1]["seed"],
+            curr_frame["compositions"][2]["prompt"],
+            curr_frame["compositions"][2]["seed"],
+            curr_frame["compositions"][3]["prompt"],
+            curr_frame["compositions"][3]["seed"],
         )
     except Exception as e:
-        return (
-            gr.update(choices=[]),
-            gr.update(choices=[]),
-            {},
-            {},
-            {},
-            gr.Checkbox(visible=False),
-        )
+        print("Error in load_data:", str(e))  # Debugging the error
+        gr.Warning("Failed to load data. Check logs!")
+
+
+def update_characters(character_data: dict, current_frame: int, frame_hash_map: dict, episode_data: dict):
+    scene_num, frame_num = (
+        frame_hash_map[current_frame]["scene"],
+        frame_hash_map[current_frame]["frame"],
+    )
+    curr_frame = episode_data["scenes"][scene_num]["frames"][frame_num]
+    return gr.CheckboxGroup(
+        choices=list(character_data.keys()),
+        value=[char["name"] for char in curr_frame["characters"]],
+    )
 
 
 def load_data_next(
-    episodes_data: list,
+    comic_id: str,
     current_episode: int,
     current_frame: int,
-    is_developer: bool,
-    character_data,
+    frame_hash_map: dict,
+    episode_data: dict,
 ):
-    if current_frame + 1 < len(episodes_data[current_episode]):
+    if current_frame + 1 < list(frame_hash_map.keys())[-1]:
         current_frame += 1
-    elif current_episode + 1 < len(episodes_data):
-        current_episode += 1
-        current_frame = 0
     else:
-        return [], current_episode, current_frame
+        current_episode += 1
+        episode_data, frame_hash_map = load_episode_data(comic_id, current_episode)
+        if len(episode_data) < 1:
+            gr.Warning("All episodes finished.")
+            return
+        current_frame = min(list(frame_hash_map.keys()))
     return (
         gr.update(value=current_episode),
-        gr.update(value=current_frame),
-        *load_data_inner(
-            episodes_data, current_episode, current_frame, is_developer, character_data
-        ),
+        gr.update(choices=list(frame_hash_map.keys()), value=current_frame),
+        current_episode,
+        current_frame,
+        episode_data,
+        frame_hash_map,
     )
 
 
 def load_data_prev(
-    episodes_data: list,
+    comic_id: str,
     current_episode: int,
     current_frame: int,
-    is_developer: bool,
-    character_data,
+    frame_hash_map: dict,
+    episode_data: dict,
 ):
-    if current_frame - 1 >= 0:
+    if current_frame - 1 >= list(frame_hash_map.keys())[0]:
         current_frame -= 1
-    elif current_episode - 1 > min(list(episodes_data.keys())):
-        current_episode -= 1
-        current_frame = 0
     else:
-        return [], current_episode, current_frame
+        current_episode -= 1
+        episode_data, frame_hash_map = load_episode_data(comic_id, current_episode)
+        if len(episode_data) < 1:
+            gr.Warning("No previous episode found.")
+            return
+        current_frame = min(list(frame_hash_map.keys()))
     return (
         gr.update(value=current_episode),
-        gr.update(value=current_frame),
-        *load_data_inner(
-            episodes_data, current_episode, current_frame, is_developer, character_data
-        ),
+        gr.update(choices=list(frame_hash_map.keys()), value=current_frame),
+        current_episode,
+        current_frame,
+        episode_data,
+        frame_hash_map,
     )
 
 
-def load_from_dropdown(
-    episodes_data: dict,
-    selected_episode: int,
-    selected_frame: int,
-    is_developer: bool,
-    character_data,
+def save_image(
+    selected_image: ..., comic_id: str, current_episode: int, current_frame: int
 ):
-    return (
-        gr.update(value=selected_episode),
-        gr.update(value=selected_frame),
-        *load_data_inner(
-            episodes_data,
-            selected_episode,
-            selected_frame,
-            is_developer,
-            character_data,
-        ),
-        # update_characters(character_data)
-    )
-
-
-def get_scene_number(details, current_episode, current_frame):
-    episode_details = details.get(current_episode)
-    if not episode_details:
-        print(f"Episode {current_episode} not found!")
-        return None
-
-    # Determine scene number and frame number within the scene
-    scene_num, frame_num_in_scene = None, 0
-    for scene_idx, cumulative_frame_count in enumerate(episode_details.items()):
-        if current_frame < cumulative_frame_count[1]:
-            scene_num = cumulative_frame_count[0]
-            frame_num_in_scene = current_frame - (episode_details.get(scene_num - 1, 0))
-            break
-    return scene_num, frame_num_in_scene
-
-
-def load_dropdown_fn(selected_episode):
-    return (gr.update(value=selected_episode), gr.update(value=0), selected_episode, 0)
-
-
-def load_dropdown_fn_v2(selected_frame):
-    return selected_frame
-
-
-def save_image(selected_image, comic_id: str, current_episode: int, current_frame: int):
     with Image.open(selected_image[0]) as img:
         img_bytes = io.BytesIO()
         img.convert("RGB").save(img_bytes, "JPEG")
@@ -293,224 +231,201 @@ def save_image(selected_image, comic_id: str, current_episode: int, current_fram
     gr.Info("Saved Image successfully!")
 
 
-def toggle_developer_options(
-    is_developer: bool, prompt_1, prompt_2, prompt_3, prompt_4, setting
-):
-    if is_developer:
-        # Return visibility updates for the developer options along with the values
-        return gr.update(visible=True), prompt_1, prompt_2, prompt_3, prompt_4, setting
-    else:
-        # Hide the developer options and return only the updated visibility
-        return gr.update(visible=False), prompt_1, prompt_2, prompt_3, prompt_4, setting
-
-
-def regenerate_composition_data(
-    image_description,
-    narration,
-    character,
-    dialouge,
-    location,
-    setting,
-    chars,
+def regenerate_compositions(
+    image_description: str,
+    narration: str,
+    character: str,
+    dialouge: str,
+    location: str,
+    setting: str,
+    rel_chars: list,
     current_episode: int,
     current_frame: int,
     episodes_data: dict,
+    frame_hash_map: dict,
+    character_data: dict,
 ):
     try:
         print(
             f"Regenerating composition data for episode {current_episode}, frame {current_frame}"
         )
-        frame = episodes_data[current_episode][current_frame]
-
-        try:
-            print("Creating prompt template for composition generation")
-            prompt_template = jinja2.Template(
-                script_gen.generate_image_compositions_user_prompt
-            )
-        except Exception as e:
-            print(f"Error creating prompt template: {e}")
-            raise
+        scene_num, frame_num = (
+            frame_hash_map[current_frame]["scene"],
+            frame_hash_map[current_frame]["frame"],
+        )
+        prev_frame = {}
+        if frame_num-1 > 0:
+            prev_frame = episodes_data["scenes"][scene_num]["frames"][frame_num-1]
 
         try:
-            print("Rendering prompt with frame details")
+            related_chars = [character_data[char] for char in rel_chars]
             prompt_dict = {
                 "system": script_gen.generate_image_compositions_instruction,
-                "user": prompt_template.render(
+                "user": jinja2.Template(
+                    script_gen.generate_image_compositions_user_prompt
+                ).render(
                     {
                         "FRAME": {
                             "description": image_description,
                             "narration": narration,
-                            "character_dilouge": dialouge,
-                            "character": character,
+                            "audio_cue_text": dialouge,
+                            "audio_cue_character": character,
                             "location": location,
-                            "setting": setting,
-                            "all_characters": chars,
-                        }
+                            "frame_setting": setting,
+                            "characters": json.dumps(related_chars),
+                        },
+                        "LOCATION_DESCRIPTION": prev_frame.get("location", ""),
+                        "frame_settings": prev_frame.get("frame_setting", ""),
                     }
                 ),
             }
-        except Exception as e:
-            print(f"Error rendering prompt: {e}")
-            raise
 
-        try:
             print("Generating compositions using LLM")
             compositions = llm.generate_valid_json_response(prompt_dict)
-        except Exception as e:
-            print(f"Error generating compositions: {e}")
-            raise
-
-        try:
-            print("Updating frame compositions")
-            frame.compositions = [
-                Composition(
-                    **comp,
-                    seed=(
-                        frame.compositions[idx].seed
-                        if idx < len(frame.compositions)
-                        else ""
-                    ),
-                    image=(
-                        frame.compositions[idx].image
-                        if idx < len(frame.compositions)
-                        else ""
-                    ),
-                )
-                for idx, comp in enumerate(compositions["compositions"])
-            ]
+            comps = compositions["compositions"]
         except Exception as e:
             print(f"Error updating frame compositions: {e}")
             raise
 
         print("Composition data regenerated successfully.")
         return [
-            frame.compositions[0].prompt,
-            frame.compositions[0].seed,
-            frame.compositions[1].prompt,
-            frame.compositions[1].seed,
-            frame.compositions[2].prompt,
-            frame.compositions[2].seed,
-            frame.compositions[3].prompt,
-            frame.compositions[3].seed,
+            comps[0]["prompt"],
+            comps[1]["prompt"],
+            comps[2]["prompt"],
+            comps[3]["prompt"],
         ]
     except Exception as e:
         print(f"Error in regenerate_composition_data: {e}")
         return [""] * 8
 
 
-def regenerate_data(
-    comic_id,
-    current_episode,
-    current_frame,
-    episodes_data,
-    character_data,
-    visual_style,
-    height,
-    width,
-    details,
-    chars,
+def regenerate_images(
+    current_episode: int,
+    current_frame: int,
+    visual_style: str,
+    height: int,
+    width: int,
+    character_data: dict,
+    rel_chars: dict,
+    prompt_1: str,
+    seed_1: str,
+    prompt_2: str,
+    seed_2: str,
+    prompt_3: str,
+    seed_3: str,
+    prompt_4: str,
+    seed_4: str,
 ):
-    images = []
-    image_data_b64 = []
-    print(chars)
-    try:
-        current_scene, frame_num_in_scene = get_scene_number(
-            details, current_episode, current_frame
-        )
+    image_list = []
+    try:        
         print(
-            f"Regenerating data for episode {current_episode}, Scene: {current_scene}, frame {frame_num_in_scene}"
+            f"Regenerating data for episode {current_episode}, and frame {current_frame}"
         )
-        frame = episodes_data[current_episode][current_frame]
-        related_chars = [character_data[ch] for ch in chars]
+        related_chars = [character_data[ch]["profile_image"] for ch in rel_chars]
+        new_compositions = [
+            {
+                "prompt": prompt_1,
+                "seed": seed_1,
+            },
+            {
+                "prompt": prompt_2,
+                "seed": seed_2,
+            },
+            {
+                "prompt": prompt_3,
+                "seed": seed_3,
+            },
+            {
+                "prompt": prompt_4,
+                "seed": seed_4,
+            },
+        ]
 
-        for i, composition in enumerate(frame.compositions):
+        for i, composition in enumerate(new_compositions):
             try:
                 print(f"Generating image for composition {i}")
+                prompt = composition["prompt"]
+                if "NOCHAR" in prompt:
+                    prompt = prompt.replace(
+                        "NOCHAR", ""
+                    )
                 payload = {
-                    "prompt": composition.prompt,
+                    "prompt": prompt,
                     "characters": related_chars,
                     "parameters": {
                         "height": height,
                         "width": width,
                         "visual_style": visual_style,
-                        "seed": composition.seed,
+                        "seed": composition["seed"],
                     },
                 }
 
-                try:
-                    print(f"Sending request to generate image for composition {i}")
-                    data = iowrapper.get_valid_post_response(
-                        url=f"{parameters.MODEL_SERVER_URL}generate_image",
-                        payload=payload,
-                    )
-                    print(f"Image generated for composition {i}. Decoding image data.")
-                    image_data = io.BytesIO(base64.b64decode(data["image"]))
-                    image_data_b64.append(data["image"])
-                    images.append(Image.open(image_data))
-                except Exception as e:
-                    print(f"Error generating image for composition {i}: {e}")
-                    continue
+                data = iowrapper.get_valid_post_response(
+                    url=f"{parameters.MODEL_SERVER_URL}generate_image",
+                    payload=payload,
+                )
+                image_list.append(Image.open(io.BytesIO(base64.b64decode(data["image"]))))
             except Exception as e:
                 print(f"Error processing composition {i}: {e}")
                 continue
 
-        print("Data regeneration completed.")
-        return images, image_data_b64
+        print(f"Generated new images for episode: {current_episode} and frame: {current_frame}")
+        print(f"Length of image list: {len(image_list)}")
+        return image_list
     except Exception as e:
         print(f"Error in regenerate_data: {e}")
-        return [], []
+        gr.Warning("Failed to generate new images!")
+        return []
 
 
-def save_image_compositions(
+def save_comic_data(
     current_episode: int,
     current_frame: int,
-    details: dict,
+    episode_data: dict,
     comic_id: str,
-    image_description,
-    narration,
-    character,
-    dialogue,  # Fixed typo from 'dialouge' to 'dialogue'
-    location,
-    setting,
-    chars,
-    prompt_1,
-    prompt_2,
-    prompt_3,
-    prompt_4,
+    image_description: str,
+    narration: str,
+    character: str,
+    dialogue: str,
+    location: str,
+    setting: str,
+    prompt_1: str,
+    prompt_2: str,
+    prompt_3: str,
+    prompt_4: str,
+    frame_hash_map: dict,
+    rel_chars: list,
+    character_data: dict,
+    images: list
 ):
     try:
-        current_scene, frame_num_in_scene = get_scene_number(
-            details, current_episode, current_frame
+        scene_num, frame_num = (
+            frame_hash_map[current_frame]["scene"],
+            frame_hash_map[current_frame]["frame"],
         )
+        curr_frame = episode_data["scenes"][scene_num]["frames"][frame_num]
         print(
-            f"Saving image components for episode {current_episode}, frame {frame_num_in_scene}"
-        )
-
-        # Fetch episode data from S3
-        episode_path = f"s3://blix-demo-v0/{comic_id}/episodes/episode-{current_episode}/episode.json"
-        print(f"Fetching episode from S3: {episode_path}")
-        episode_json = aws_utils.fetch_from_s3(episode_path).decode("utf-8")
-        episode = json.loads(episode_json)
-
-        frame_data = episode["scenes"][current_scene]["frames"][frame_num_in_scene]
-        print(
-            f"Updating compositions for scene {current_scene}, frame {frame_num_in_scene}"
+            f"Saving comic data for episode {current_episode}, frame {frame_num}"
         )
 
         # Update compositions with prompts
         prompts_list = [prompt_1, prompt_2, prompt_3, prompt_4]
-        frame_data["compositions"] = [
-            {
-                "prompt": prompts_list[i],
-                "shot_type": comp["shot_type"],
-                "seed": comp["seed"],
-                "image": comp["image"],
-            }
-            for i, comp in enumerate(frame_data["compositions"])
-        ]
+        for i, comp in enumerate(curr_frame["compositions"]):
+            comp["prompt"] = prompts_list[i]
+            # Save new images to S3
+            with Image.open(images[i][0]) as img:
+                img_bytes = io.BytesIO()
+                img.convert("RGB").save(img_bytes, "JPEG")
+                img_bytes.seek(0)
+                aws_utils.save_to_s3(
+                    parameters.AWS_BUCKET,
+                    f"{comic_id}/episodes/episode-{current_episode}/compositions/scene-{scene_num}/frame-{frame_num}",
+                    img_bytes,
+                    f"{i}.jpg",
+                )
 
-        # Batch update frame data
-        frame_data.update(
+        # Update frame data
+        curr_frame.update(
             {
                 "description": image_description,
                 "narration": narration,
@@ -518,48 +433,19 @@ def save_image_compositions(
                 "location": location,
                 "setting": setting,
                 "audio_cue_character": character,
+                "characters": [character_data[char] for char in rel_chars],
             }
         )
 
         # Save the updated episode back to S3
-        print(f"Saving updated episode to S3 at {episode_path}")
+        print(f"Saving updated episode {current_episode} to S3")
         aws_utils.save_to_s3(
             bucket_name=parameters.AWS_BUCKET,
             folder_name=f"{comic_id}/episodes/episode-{current_episode}",
-            content=json.dumps(episode),
+            content=episode_data,
             file_name="episode.json",
         )
-
-        gr.Info("Compositions saved successfully!")
-        print(current_scene)
-        return current_scene
-
-    except Exception as e:
-        print(f"Error in save_image_compositions: {e}")
-        return None
-
-
-def save_images(image_data_b64, current_episode, current_frame, comic_id, details):
-    try:
-        current_scene, frame_num_in_scene = get_scene_number(
-            details, current_episode, current_frame
-        )
-        print(
-            f"Saving images for episode {current_episode}, Scene {current_scene} ,frame {frame_num_in_scene}."
-        )
-        for i, image_data in enumerate(image_data_b64):
-            try:
-                print(f"Saving image {i} to S3")
-                aws_utils.save_to_s3(
-                    parameters.AWS_BUCKET,
-                    f"{comic_id}/episodes/episode-{current_episode}/compositions/scene-{current_scene}/frame-{frame_num_in_scene}",
-                    io.BytesIO(base64.b64decode(image_data)),
-                    f"{i}.jpg",
-                )
-            except Exception as e:
-                print(f"Error saving image {i} to S3: {e}")
-                continue
-        print("All images saved successfully!")
-        gr.Info("All images saved successfully!")
+        gr.Info("Comic data saved successfully!")
     except Exception as e:
-        print(f"Error in save_images: {e}")
+        print(f"Error in saving comic data: {e}")
+        gr.Warning("Failed to save data for the comic!")

parameters.py

@@ -1,7 +1,7 @@
 import os
-# from dotenv import load_dotenv
+from dotenv import load_dotenv
 
-# load_dotenv()
+load_dotenv()
 
 
 AWS_BUCKET = os.getenv("AWS_BUCKET")
@@ -11,4 +11,6 @@ os.environ["S3_BUCKET_NAME"] = os.getenv("AWS_BUCKET")
 VISUAL_CHOICES = ["DARK", "FLUX_COMIC", "GHIBLI_COMIC"]
 MAX_TRIES = os.getenv("MAX_TRIES")
 OPEN_AI_API_KEY = os.getenv("OPEN_AI_KEY")
-MODEL_SERVER_URL = os.getenv("MODEL_SERVER_URL")
+MODEL_SERVER_URL = os.getenv("MODEL_SERVER_URL")
+IMG_HEIGHT=os.getenv("IMG_HEIGHT")
+IMG_WIDTH=os.getenv("IMG_WIDTH")

script_gen.py

@@ -1,53 +1,60 @@
 generate_image_compositions_instruction = """\
-As a visual artist and cinematographer with extensive knowledge of different camera angles, visual styles, and aesthetics, your task is to analyze the provided frame details and create four distinct compositions. Each composition should follow a specific narrative structure in its prompt construction:
+As a visual artist and cinematographer with extensive knowledge of different camera angles, visual styles, and aesthetics, your task is to analyze the provided frame details and create four distinct compositions, maintaining consistency with previous frame's location and settings. Each composition should follow a specific narrative structure in its prompt construction:
 
 1) Output Requirements in json:
 {
     "compositions": [
         {
             "prompt": "Detailed visual description following the structure below (max 77 tokens)",
-            "shot_type": "Optimal cinematographic shot"
+            "shot_type": "Optimal cinematographic shot",
         }
     ]
 }
 
 2) Prompt Structure (in this specific order):
    a) Begin with the environment and setting:
-      - Establish the broader landscape/location first
-      - Describe key environmental elements
-      - Set the atmospheric conditions
-      - Define the lighting and mood
-   
+      - Establish the broader landscape/location first, maintaining consistency with previous frame
+      - Describe key environmental elements that persist from previous frame
+      - Set the atmospheric conditions aligned with previous settings
+      - Define the lighting and mood that follows from previous frame
+
    b) Then layer in the scene elements:
       - How different parts of the environment interact
       - Spatial relationships and depth
-      - Textures and materials
+      - Textures and materials consistent with previous frame
       - Any dynamic elements (movement, weather effects)
-   
+
    c) Finally, integrate characters (if applicable):
       - Their position within the established environment
       - How they interact with the space
       - Their expressions and actions as part of the scene
-      
+
 3) Each composition should:
    - Flow naturally like a single, cohesive description
    - Prioritize environmental storytelling
    - Build the scene progressively from background to foreground
    - Maintain consistent atmosphere throughout
+   - Ensure visual continuity with previous frame's location and settings
 
 4) For NO-CHAR compositions:
    Focus entirely on a and b of the prompt structure, with extra emphasis on:
-   - Environmental details and patterns
-   - Architectural elements
+   - Environmental details and patterns that match previous frame
+   - Architectural elements maintaining previous frame's style
    - Natural phenomena
-   - Atmospheric qualities
+   - Atmospheric qualities aligned with previous settings
 
-Note: Avoid jumping between environment and character descriptions. Each element should flow naturally into the next, creating a unified visual narrative.
+Note: Avoid jumping between environment and character descriptions. Each element should flow naturally into the next, creating a unified visual narrative while maintaining consistency with previous frame's location and settings.\
 """
 
 generate_image_compositions_user_prompt = """\
 Here's are the details:
 
-## Synopsis:
+## Current Frame:
 {{FRAME}}
+
+Previous Frame location:
+{{LOCATION_DESCRIPTION}}
+
+Previous Frame settings:
+{{frame_settings}}
 """