comic grading updated code

app.py

@@ -1,4 +1,5 @@
 """Driver File."""
+
 import gradio as gr
 import core
 import parameters
@@ -14,6 +15,7 @@ with gr.Blocks() as demo:
     character_data = gr.State({})
     current_frame_data = gr.State(None)
     details = gr.State({})
+    image_data_b64 = gr.State([])
 
     with gr.Row():
         with gr.Column():
@@ -74,33 +76,43 @@ with gr.Blocks() as demo:
                     gr.Markdown("## Composition #4")
                     prompt_4 = gr.TextArea(label="Image Prompt")
                     seed_4 = gr.Textbox(label="Generation Seed")
-            regenerate_comps_btn = gr.Button(value="Regenerate Compositions")
+
+            with gr.Row():
+                regenerate_comps_btn = gr.Button(value="Regenerate Compositions")
+                save_components_btn = gr.Button(value="Save Components")
 
         with gr.Column():
-            negative_prompt = gr.TextArea(value="", label="Negative Prompt", )
+            negative_prompt = gr.TextArea(
+                value="",
+                label="Negative Prompt",
+            )
             chars = gr.Textbox(value="[]", label="Related Characters")
             with gr.Row():
                 height = gr.Textbox(value="1024", label="Image Height")
                 width = gr.Textbox(value="1024", label="Image Width")
                 visual_style = gr.Dropdown(
-                    choices=parameters.VISUAL_CHOICES, label="Current Visual Style", interactive=True
+                    choices=parameters.VISUAL_CHOICES,
+                    label="Current Visual Style",
+                    interactive=True,
                 )
-        
+
         with gr.Row():
             regenerate_btn = gr.Button("Regenerate")
             save_btn = gr.Button("Save")
 
-############################################ EVENTS ############################################
+    ############################################ EVENTS ############################################
     load_metadata.click(
         core.load_metadata_fn,
         inputs=[comic_id],
         outputs=[
             episode_dropdown,
             frame_dropdown,
+            current_episode,
+            current_frame,
             episodes_data,
             character_data,
             details,
-            developer
+            developer,
         ],
     )
 
@@ -229,30 +241,43 @@ with gr.Blocks() as demo:
 
     developer.change(
         core.toggle_developer_options,
-        inputs=[developer],
-        outputs=[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,
         inputs=[
+            image_description,
+            narration,
+            character,
+            dialouge,
+            location,
+            setting,
+            chars,
             current_episode,
-            current_scene,
             current_frame,
             episodes_data,
-            details,
-            comic_id
         ],
         outputs=[
             prompt_1,
             seed_1,
-            prompt_2, 
+            prompt_2,
             seed_2,
             prompt_3,
             seed_3,
             prompt_4,
-            seed_4
-        ]
+            seed_4,
+        ],
     )
     regenerate_btn.click(
         core.regenerate_data,
@@ -265,10 +290,42 @@ with gr.Blocks() as demo:
             character_data,
             visual_style,
             height,
-            width
+            width,
         ],
-        outputs=[images]
+        outputs=[images, image_data_b64],
     )
 
-demo.launch(auth=("admin", "Qrt@12*34#immersfy"), share=True, ssr_mode=False,  debug=True)
-# demo.launch(share=True, ssr_mode=False,  debug=True)
+    save_components_btn.click(
+        core.save_image_compositions,
+        inputs=[
+            current_episode,
+            current_frame,
+            details,
+            comic_id,
+            image_description,
+            narration,
+            character,
+            dialouge,
+            location,
+            setting,
+            chars,
+            prompt_1,
+            prompt_2,
+            prompt_3,
+            prompt_4,
+        ],
+        outputs=[current_scene],
+    )
+
+    save_btn.click(
+        core.save_images,
+        inputs=[
+            image_data_b64,
+            current_episode,
+            current_frame,
+            current_scene,
+            comic_id,
+        ],
+    )
+# 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

@@ -158,6 +158,8 @@ def load_metadata_fn(comic_id: str):
             gr.update(
                 choices=range(len(episodes_data[current_episode])), value=current_frame
             ),
+            current_episode,
+            current_frame,
             episodes_data,
             character_data,
             details,
@@ -240,91 +242,96 @@ 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):
+def toggle_developer_options(
+    is_developer: bool, prompt_1, prompt_2, prompt_3, prompt_4, setting
+):
     if is_developer:
-        return gr.Column(visible=True)
-    gr.Column(visible=False)
+        # 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,
     current_episode: int,
-    current_scene: int,
     current_frame: int,
     episodes_data: dict,
-    details: dict,
-    comic_id: str,
 ):
     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
+        print(
+            f"Regenerating composition data for episode {current_episode}, frame {current_frame}"
+        )
+        frame = episodes_data[current_episode][current_frame]
 
-        if scene_num is None:
-            return
+        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
 
-        frame = episodes_data[current_episode][current_frame]
-        prompt_template = jinja2.Template(
-            script_gen.generate_image_compositions_user_prompt
-        )
-        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,
+        try:
+            print("Rendering prompt with frame details")
+            prompt_dict = {
+                "system": script_gen.generate_image_compositions_instruction,
+                "user": prompt_template.render(
+                    {
+                        "FRAME": {
+                            "description": image_description,
+                            "narration": narration,
+                            "character_dilouge": dialouge,
+                            "character": character,
+                            "location": location,
+                            "setting": setting,
+                            "all_characters": chars,
+                        }
                     }
-                }
-            ),
-        }
-
-        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"])
-        ]
+            }
+        except Exception as e:
+            print(f"Error rendering prompt: {e}")
+            raise
 
-        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
-        ]
+        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
 
-        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",
-        )
+        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"])
+            ]
+        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,
@@ -336,6 +343,7 @@ def regenerate_composition_data(
             frame.compositions[3].seed,
         ]
     except Exception as e:
+        print(f"Error in regenerate_composition_data: {e}")
         return [""] * 8
 
 
@@ -351,35 +359,169 @@ def regenerate_data(
     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,
-            "characters": related_chars,
-            "parameters": {
-                "height": height,
-                "width": width,
-                "visual_style": visual_style,
-                "seed": composition.seed,
-            },
-        }
+    image_data_b64 = []
+    try:
+        print(
+            f"Regenerating data for episode {current_episode}, scene {current_scene}, frame {current_frame}"
+        )
+        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):
+            try:
+                print(f"Generating image for composition {i}")
+                payload = {
+                    "prompt": composition.prompt,
+                    "characters": related_chars,
+                    "parameters": {
+                        "height": height,
+                        "width": width,
+                        "visual_style": visual_style,
+                        "seed": composition.seed,
+                    },
+                }
 
-        try:
-            data = iowrapper.get_valid_post_response(
-                url=f"{parameters.MODEL_SERVER_URL}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
+                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(image_data)
+                    images.append(Image.open(image_data))
+                except Exception as e:
+                    print(f"Error generating image for composition {i}: {e}")
+                    continue
+            except Exception as e:
+                print(f"Error processing composition {i}: {e}")
+                continue
+
+        print("Data regeneration completed.")
+        return images, image_data_b64
+    except Exception as e:
+        print(f"Error in regenerate_data: {e}")
+        return [], []
+
+
+def save_image_compositions(
+    current_episode: int,
+    current_frame: int,
+    details: 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,
+):
+    try:
+        print(
+            f"Saving image components for episode {current_episode}, frame {current_frame}"
+        )
 
-    return images
+        # Fetch episode details early and return if not found
+        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
+
+        if scene_num is None:
+            print(f"Scene not found for frame {current_frame}.")
+            return None
+
+        # 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"][scene_num]["frames"][frame_num_in_scene]
+        print(
+            f"Updating compositions for scene {scene_num}, frame {frame_num_in_scene}"
+        )
+
+        # 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"])
+        ]
+
+        # Batch update frame data
+        frame_data.update(
+            {
+                "description": image_description,
+                "narration": narration,
+                "audio_cue_text": dialogue,
+                "location": location,
+                "setting": setting,
+                "audio_cue_character": character,
+            }
+        )
+
+        # Save the updated episode back to S3
+        print(f"Saving updated episode to S3 at {episode_path}")
+        aws_utils.save_to_s3(
+            bucket_name=parameters.AWS_BUCKET,
+            folder_name=f"{comic_id}/episodes/episode-{current_episode}",
+            content=json.dumps(episode),
+            file_name="episode.json",
+        )
+
+        gr.Info("Components saved successfully!")
+        return scene_num
+
+    except Exception as e:
+        print(f"Error in save_image_compositions: {e}")
+        return None
+
+
+def save_images(
+    image_data_b64,
+    current_episode,
+    current_frame,
+    current_scene,
+    comic_id,
+):
+    try:
+        print(
+            f"Saving images for scene {current_scene}, episode {current_episode}, frame {current_frame}."
+        )
+        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-{current_frame}",
+                    image_data,
+                    f"{i}.jpg",
+                )
+            except Exception as e:
+                print(f"Error saving image {i} to S3: {e}")
+                continue
+        gr.Info("All Images saved successfully!")
+    except Exception as e:
+        print(f"Error in save_images: {e}")