regeneratio added, WIP

app.py

@@ -5,11 +5,13 @@ import parameters
 
 ############################################ LAYOUT ############################################
 with gr.Blocks() as demo:
+    # ... [previous state variables remain the same] ...
     selected_image = gr.State(None)
     current_episode = gr.State(-1)
     current_scene = gr.State(-1)
     current_frame = gr.State(-1)
     episodes_data = gr.State({})
+    character_data = gr.State({})
     current_frame_data = gr.State(None)
 
     with gr.Row():
@@ -57,23 +59,19 @@ with gr.Blocks() as demo:
                 with gr.Column():
                     gr.Markdown("## Composition #1")
                     prompt_1 = gr.TextArea(label="Image Prompt")
-                    # shot_1 = gr.Textbox(label="Shot Type")
                     seed_1 = gr.Textbox(label="Generation Seed")
                 with gr.Column():
                     gr.Markdown("## Composition #2")
                     prompt_2 = gr.TextArea(label="Image Prompt")
-                    # shot_2 = gr.Textbox(label="Shot Type")
                     seed_2 = gr.Textbox(label="Generation Seed")
             with gr.Row():
                 with gr.Column():
                     gr.Markdown("## Composition #3")
                     prompt_3 = gr.TextArea(label="Image Prompt")
-                    # shot_3 = gr.Textbox(label="Shot Type")
                     seed_3 = gr.Textbox(label="Generation Seed")
                 with gr.Column():
                     gr.Markdown("## Composition #4")
                     prompt_4 = gr.TextArea(label="Image Prompt")
-                    # shot_4 = gr.Textbox(label="Shot Type")
                     seed_4 = gr.Textbox(label="Generation Seed")
             regenerate_comps_btn = gr.Button(value="Regenerate Compositions")
 
@@ -99,7 +97,8 @@ with gr.Blocks() as demo:
             episode_dropdown,
             frame_dropdown,
             episodes_data,
-            developer,
+            character_data,
+            developer
         ],
     )
 
@@ -239,12 +238,31 @@ with gr.Blocks() as demo:
             current_scene,
             current_frame,
             episodes_data
-            ],
-        outputs=[]
+        ],
+        outputs=[
+            prompt_1,
+            seed_1,
+            prompt_2, 
+            seed_2,
+            prompt_3,
+            seed_3,
+            prompt_4,
+            seed_4
+        ]
     )
     regenerate_btn.click(
         core.regenerate_data,
-        inputs=[],
+        inputs=[
+            comic_id,
+            current_episode,
+            current_scene,
+            current_frame,
+            episodes_data,
+            character_data,
+            visual_style,
+            height,
+            width
+        ],
         outputs=[]
     )
 

core.py

@@ -10,7 +10,7 @@ import base64
 import aws_utils
 import parameters
 import script_gen
-import io as iowrapper
+import inout as iowrapper
 import openai_wrapper
 
 AWS_BUCKET = parameters.AWS_BUCKET
@@ -57,11 +57,11 @@ def load_data_inner(
 ):
     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,
@@ -91,6 +91,12 @@ def load_metadata_fn(comic_id: str):
     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"
@@ -124,6 +130,7 @@ def load_metadata_fn(comic_id: str):
             choices=range(len(episodes_data[current_episode])), value=current_frame
         ),
         episodes_data,
+        character_data,
         gr.Checkbox(visible=True),
     )
 
@@ -223,11 +230,10 @@ def regenerate_composition_data(
     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]
-    print(frame)
-    
+
     # Generate the prompt for compositions
     prompt_dict = {
         "system": script_gen.generate_image_compositions_instruction,
@@ -235,55 +241,89 @@ def regenerate_composition_data(
             source=script_gen.generate_image_compositions_user_prompt
         ).render(
             {
-                "FRAME": frame,
+                "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 compositions using LLM
+    # Generate composition  s using LLM
     compositions = llm.generate_valid_json_response(prompt_dict)
-    print(compositions)
-
     # Update frame with new compositions
     frame.compositions = [
         Composition(
             **composition,
-            seed="",  # Set default empty seed
-            image=""  # Set default empty image
-        )  # Create a Composition dataclass from the response
-        for composition in compositions["compositions"]
+            seed=frame.compositions[idx].seed if idx < len(frame.compositions) else "",
+            image=(
+                frame.compositions[idx].image if idx < len(frame.compositions) else ""
+            ),
+        )
+        for idx, composition in enumerate(compositions["compositions"])
     ]
-    
+
     # 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.")
+    print(
+        f"Updated frame {current_frame} in episode {current_episode} with new compositions."
+    )
+
+    # 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,
+    ]
 
 
 def regenerate_data(
-    frame_data: ComicFrame,
+    comic_id,
+    current_episode,
+    current_scene,
+    current_frame,
+    episodes_data,
+    character_data,
+    visual_style,
+    height,
+    width,
 ):
-    pass
-#     # for
-#     payload = {
-#         "prompt": composition.prompt,
-#         "characters": related_chars,
-#         "parameters": {
-#             "height": parameters.IMG_HEIGHT,
-#             "width": parameters.IMG_WIDTH,
-#             "visual_style": visual_style,
-#             "seed": seed_val,
-#         },
-#     }
-
-#     data = iowrapper.get_valid_post_response(
-#         url=parameters.MODEL_SERVER_URL + "/generate_image",
-#         payload=payload,
-#     )
-#     image_data = io.BytesIO(base64.b64decode(data["image"]))
-#     path = aws_utils.save_to_s3(
-#         parameters.AWS_BUCKET,
-#         f"{self.id}/episodes/episode-{episode_num}/compositions/scene-{scene_num}/frame-{frame_num}",
-#         image_data,
-#         f"{num}.jpg",
-#     )
+
+    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,
+            },
+        }
+
+        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

io.py → inout.py

@@ -33,7 +33,7 @@ def get_valid_post_response(url: str, payload: dict) -> Dict[str, Any]:
     Returns:
         Dict[str, Any]: The JSON response from the request.
     """
-    for _ in range(parameters.MAX_TRIES):
+    for _ in range(int(parameters.MAX_TRIES)):
         try:
             response, status_code = send_post_request(url, payload)
             if status_code != 200:

openai_wrapper.py

@@ -67,7 +67,7 @@ class OpenAIModel(abc.ABC):
         self,
         prompt_dict: Mapping[str, str],
         max_output_tokens: int = None,
-        temperature: int = 0.6,
+        temperature: int = 0.7,
     ) -> str:
         """Generate a response with retries, returning a valid JSON."""
         for _ in range(int(parameters.MAX_TRIES)):

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")