regeneration part added

app.py

@@ -13,6 +13,7 @@ with gr.Blocks() as demo:
     episodes_data = gr.State({})
     character_data = gr.State({})
     current_frame_data = gr.State(None)
+    details = gr.State({})
 
     with gr.Row():
         with gr.Column():
@@ -98,6 +99,7 @@ with gr.Blocks() as demo:
             frame_dropdown,
             episodes_data,
             character_data,
+            details,
             developer
         ],
     )
@@ -158,7 +160,7 @@ with gr.Blocks() as demo:
 
     next_button.click(
         core.load_data_next,
-        inputs=[episodes_data, current_episode, current_frame],
+        inputs=[episodes_data, current_episode, current_frame, developer],
         outputs=[
             episode_dropdown,
             frame_dropdown,
@@ -187,7 +189,7 @@ with gr.Blocks() as demo:
 
     prev_button.click(
         core.load_data_prev,
-        inputs=[episodes_data, current_episode, current_frame],
+        inputs=[episodes_data, current_episode, current_frame, developer],
         outputs=[
             episode_dropdown,
             frame_dropdown,
@@ -237,7 +239,9 @@ with gr.Blocks() as demo:
             current_episode,
             current_scene,
             current_frame,
-            episodes_data
+            episodes_data,
+            details,
+            comic_id
         ],
         outputs=[
             prompt_1,
@@ -263,7 +267,8 @@ with gr.Blocks() as demo:
             height,
             width
         ],
-        outputs=[]
+        outputs=[images]
     )
 
 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

@@ -1,5 +1,3 @@
-"""House of all specific functions used to control data flow."""
-
 from typing import List
 from PIL import Image
 import gradio as gr
@@ -12,6 +10,8 @@ import parameters
 import script_gen
 import inout as iowrapper
 import openai_wrapper
+import json
+from dataclasses import asdict
 
 AWS_BUCKET = parameters.AWS_BUCKET
 llm = openai_wrapper.GPT_4O_MINI
@@ -31,112 +31,151 @@ class ComicFrame:
     narration: str
     character_dilouge: str
     character: str
-    location: str  # Moved up here
+    location: str
     setting: str
     all_characters: list
-    compositions: List[Composition] = dataclasses.field(
-        default_factory=list
-    )  # Keep this as the last argument
+    compositions: List[Composition] = dataclasses.field(default_factory=list)
 
 
 def list_current_dir(bucket_name: str, folder_path: str = "") -> list:
-    response = aws_utils.S3_CLIENT.list_objects_v2(
-        Bucket=bucket_name, Prefix=folder_path, Delimiter="/"
-    )
-
-    # Check if the bucket contains objects
-    folders = []
-    if "CommonPrefixes" in response:
-        for prefix in response["CommonPrefixes"]:
-            folders.append(prefix["Prefix"])
-    return folders
+    try:
+        response = aws_utils.S3_CLIENT.list_objects_v2(
+            Bucket=bucket_name, Prefix=folder_path, Delimiter="/"
+        )
+        folders = []
+        if "CommonPrefixes" in response:
+            for prefix in response["CommonPrefixes"]:
+                folders.append(prefix["Prefix"])
+        return folders
+    except Exception as e:
+        return []
 
 
 def load_data_inner(
     episodes_data: list, current_episode: int, current_frame: int, is_developer: bool
 ):
-    images = []
-    curr_frame = episodes_data[current_episode][current_frame]
-    print(episodes_data[current_episode][current_frame])
-    # Loading the 0th frame of 0th scene in 0th episode.
-    for comp in curr_frame.compositions:
-        data = aws_utils.fetch_from_s3(comp.image)
-        images.append(Image.open(io.BytesIO(data)))
-    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,
-        curr_frame.all_characters,
-    )
+    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)))
+        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,
+            curr_frame.all_characters,
+        )
+    except Exception as e:
+        return (
+            [],
+            episodes_data,
+            current_episode,
+            current_frame,
+            gr.Textbox(),
+            gr.Textbox(),
+            gr.Textbox(),
+            gr.Textbox(),
+            gr.Textbox(),
+            "",
+            "",
+            "",
+            "",
+            "",
+            "",
+            "",
+            "",
+            "",
+            [],
+        )
 
 
 def load_metadata_fn(comic_id: str):
-    # Logic to load and return images based on comic_id and episode
-    # You can replace this with actual image paths or generation logic
-    print(f"Getting episodes for comic id: {comic_id}")
-    episodes_data = {}
-    episode_idx = []
-    character_data = {}
-    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"))
-    for name, char in char_data.items():
-        character_data[name] = char["profile_image"]
-    print(character_data)
-    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 = []
-            for scene in data["scenes"]:
-                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"]
-                            ],
+    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"))
+
+        for name, char in char_data.items():
+            character_data[name] = char["profile_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
-        ),
-        episodes_data,
-        character_data,
-        gr.Checkbox(visible=True),
-    )
+                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
+            ),
+            episodes_data,
+            character_data,
+            details,
+            gr.Checkbox(visible=True),
+        )
+    except Exception as e:
+        return (
+            gr.update(choices=[]),
+            gr.update(choices=[]),
+            {},
+            {},
+            {},
+            gr.Checkbox(visible=False),
+        )
 
 
 def load_data_next(
-    episodes_data: list, current_episode: int, current_frame: int, is_developer=False
+    episodes_data: list, current_episode: int, current_frame: int, is_developer: bool
 ):
     if current_frame + 1 < len(episodes_data[current_episode]):
         current_frame += 1
@@ -153,7 +192,7 @@ def load_data_next(
 
 
 def load_data_prev(
-    episodes_data: list, current_episode: int, current_frame: int, is_developer=False
+    episodes_data: list, current_episode: int, current_frame: int, is_developer: bool
 ):
     if current_frame - 1 >= 0:
         current_frame -= 1
@@ -180,32 +219,18 @@ def load_from_dropdown(
 
 
 def load_dropdown_fn(selected_episode):
-    return (
-        gr.update(value=selected_episode),
-        gr.update(value=0),
-        selected_episode,
-        0,
-    )
+    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,
-):
-    # Implement your AWS S3 save logic here
-    # print(f"Saving image: {selected_image}")
+def save_image(selected_image, comic_id: str, current_episode: int, current_frame: int):
     with Image.open(selected_image[0]) as img:
-        # Convert and save as JPG
         img_bytes = io.BytesIO()
         img.convert("RGB").save(img_bytes, "JPEG")
         img_bytes.seek(0)
-
         aws_utils.save_to_s3(
             AWS_BUCKET,
             f"{comic_id}/episode-{current_episode}/images",
@@ -226,64 +251,92 @@ def regenerate_composition_data(
     current_scene: int,
     current_frame: int,
     episodes_data: dict,
+    details: dict,
+    comic_id: str,
 ):
-    print(
-        f"Generating compositions for episode: {current_episode} and scene: {current_scene} and frame: {current_frame}."
-    )
-
-    # Retrieve the current frame data
-    frame = episodes_data[current_episode][current_frame]
-
-    # Generate the prompt for compositions
-    prompt_dict = {
-        "system": script_gen.generate_image_compositions_instruction,
-        "user": jinja2.Template(
-            source=script_gen.generate_image_compositions_user_prompt
-        ).render(
-            {
-                "FRAME": {
-                    "description": frame.description,
-                    "narration": frame.narration,
-                    "character_dilouge": frame.character_dilouge,
-                    "character": frame.character,
-                    "location": frame.location,
-                    "setting": frame.setting,
-                    "all_characters": frame.all_characters,
-                },
-            }
-        ),
-    }
-    # Generate composition  s using LLM
-    compositions = llm.generate_valid_json_response(prompt_dict)
-    # Update frame with new compositions
-    frame.compositions = [
-        Composition(
-            **composition,
-            seed=frame.compositions[idx].seed if idx < len(frame.compositions) else "",
-            image=(
-                frame.compositions[idx].image if idx < len(frame.compositions) else ""
-            ),
+    try:
+        episode_details = details.get(current_episode)
+        if not episode_details:
+            return
+
+        scene_num, frame_num_in_scene = None, None
+        prev_frame_count = 0
+
+        for scene_idx, cumulative_frame_count in episode_details.items():
+            if current_frame < cumulative_frame_count:
+                scene_num = scene_idx
+                frame_num_in_scene = current_frame - prev_frame_count
+                break
+            prev_frame_count = cumulative_frame_count
+
+        if scene_num is None:
+            return
+
+        frame = episodes_data[current_episode][current_frame]
+        prompt_template = jinja2.Template(
+            script_gen.generate_image_compositions_user_prompt
         )
-        for idx, composition in enumerate(compositions["compositions"])
-    ]
+        prompt_dict = {
+            "system": script_gen.generate_image_compositions_instruction,
+            "user": prompt_template.render(
+                {
+                    "FRAME": {
+                        "description": frame.description,
+                        "narration": frame.narration,
+                        "character_dilouge": frame.character_dilouge,
+                        "character": frame.character,
+                        "location": frame.location,
+                        "setting": frame.setting,
+                        "all_characters": frame.all_characters,
+                    }
+                }
+            ),
+        }
 
-    # Update the episodes_data dictionary with the modified frame
-    episodes_data[current_episode][current_frame] = frame
-    print(
-        f"Updated frame {current_frame} in episode {current_episode} with new compositions."
-    )
+        compositions = llm.generate_valid_json_response(prompt_dict)
+        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"])
+        ]
+
+        episode_path = f"s3://blix-demo-v0/{comic_id}/episodes/episode-{current_episode}/episode.json"
+        episode = json.loads(aws_utils.fetch_from_s3(episode_path).decode("utf-8"))
+        episode["scenes"][scene_num]["frames"][frame_num_in_scene]["compositions"] = [
+            asdict(comp) for comp in frame.compositions
+        ]
+
+        episode_json = json.dumps(episode)
+        aws_utils.save_to_s3(
+            bucket_name=parameters.AWS_BUCKET,
+            folder_name=f"{comic_id}/episodes/episode-{current_episode}",
+            content=episode_json,
+            file_name="episode.json",
+        )
 
-    # Return the updated composition values for the UI
-    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,
-    ]
+        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,
+        ]
+    except Exception as e:
+        return [""] * 8
 
 
 def regenerate_data(
@@ -297,9 +350,10 @@ def regenerate_data(
     height,
     width,
 ):
-
+    images = []
     frame = episodes_data[current_episode][current_frame]
     related_chars = [character_data[ch] for ch in frame.all_characters]
+
     for i, composition in enumerate(frame.compositions):
         payload = {
             "prompt": composition.prompt,
@@ -312,18 +366,20 @@ def regenerate_data(
             },
         }
 
-        data = iowrapper.get_valid_post_response(
-            url="http://10.100.111.13:4389/generate_image",
-            payload=payload,
-        )
-        image_data = io.BytesIO(base64.b64decode(data["image"]))
-        path = aws_utils.save_to_s3(
-            parameters.AWS_BUCKET,
-            f"{comic_id}/episodes/episode-{current_episode}/compositions/scene-{0}/frame-{current_frame}",
-            image_data,
-            f"{i}.jpg",
-        )
-    load_data_inner(episodes_data, current_episode, current_frame, is_developer=True)
-
-
-#     pass
+        try:
+            data = iowrapper.get_valid_post_response(
+                url="http://10.100.111.13:4389/generate_image",
+                payload=payload,
+            )
+            image_data = io.BytesIO(base64.b64decode(data["image"]))
+            aws_utils.save_to_s3(
+                parameters.AWS_BUCKET,
+                f"{comic_id}/episodes/episode-{current_episode}/compositions/scene-{0}/frame-{current_frame}",
+                image_data,
+                f"{i}.jpg",
+            )
+            images.append(Image.open(image_data))
+        except Exception as e:
+            continue
+
+    return images