Completed regeneration flow

app.py

@@ -279,21 +279,6 @@ with gr.Blocks() as demo:
             seed_4,
         ],
     )
-    regenerate_btn.click(
-        core.regenerate_data,
-        inputs=[
-            comic_id,
-            current_episode,
-            current_scene,
-            current_frame,
-            episodes_data,
-            character_data,
-            visual_style,
-            height,
-            width,
-        ],
-        outputs=[images, image_data_b64],
-    )
 
     save_components_btn.click(
         core.save_image_compositions,
@@ -316,6 +301,21 @@ with gr.Blocks() as demo:
         ],
         outputs=[current_scene],
     )
+    regenerate_btn.click(
+        core.regenerate_data,
+        inputs=[
+            comic_id,
+            current_episode,
+            current_frame,
+            episodes_data,
+            character_data,
+            visual_style,
+            height,
+            width,
+            details
+        ],
+        outputs=[images, image_data_b64],
+    )
 
     save_btn.click(
         core.save_images,
@@ -323,8 +323,8 @@ with gr.Blocks() as demo:
             image_data_b64,
             current_episode,
             current_frame,
-            current_scene,
             comic_id,
+            details
         ],
     )
 demo.launch(auth=("admin", "Qrt@12*34#immersfy"), share=True, ssr_mode=False,  debug=True)

core.py

@@ -12,6 +12,8 @@ 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
@@ -220,6 +222,22 @@ def load_from_dropdown(
     )
 
 
+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)
 
@@ -350,19 +368,22 @@ def regenerate_composition_data(
 def regenerate_data(
     comic_id,
     current_episode,
-    current_scene,
     current_frame,
     episodes_data,
     character_data,
     visual_style,
     height,
     width,
+    details,
 ):
     images = []
     image_data_b64 = []
     try:
+        current_scene, frame_num_in_scene = get_scene_number(
+            details, current_episode, current_frame
+        )
         print(
-            f"Regenerating data for episode {current_episode}, scene {current_scene}, frame {current_frame}"
+            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]
@@ -389,7 +410,7 @@ def regenerate_data(
                     )
                     print(f"Image generated for composition {i}. Decoding image data.")
                     image_data = io.BytesIO(base64.b64decode(data["image"]))
-                    image_data_b64.append(image_data)
+                    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}")
@@ -423,39 +444,22 @@ def save_image_compositions(
     prompt_4,
 ):
     try:
+        current_scene, frame_num_in_scene = get_scene_number(
+            details, current_episode, current_frame
+        )
         print(
-            f"Saving image components for episode {current_episode}, frame {current_frame}"
+            f"Saving image components for episode {current_episode}, frame {frame_num_in_scene}"
         )
 
-        # 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]
+        frame_data = episode["scenes"][current_scene]["frames"][frame_num_in_scene]
         print(
-            f"Updating compositions for scene {scene_num}, frame {frame_num_in_scene}"
+            f"Updating compositions for scene {current_scene}, frame {frame_num_in_scene}"
         )
 
         # Update compositions with prompts
@@ -492,36 +496,34 @@ def save_image_compositions(
         )
 
         gr.Info("Components saved successfully!")
-        return scene_num
+        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,
-    current_scene,
-    comic_id,
-):
+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 scene {current_scene}, episode {current_episode}, frame {current_frame}."
+            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-{current_frame}",
-                    image_data,
+                    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
-        gr.Info("All Images saved successfully!")
+        print("All images saved successfully!")
     except Exception as e:
         print(f"Error in save_images: {e}")