First commit

.gitignore

@@ -0,0 +1,3 @@
+*.env
+*.ipynb
+__pycache__/

app.py

@@ -0,0 +1,308 @@
+import textwrap
+import jinja2
+import gradio as gr
+
+from src import script_gen
+import pipeline
+
+GENRE_CHOICES = ["ROMANTIC", "HORROR", "FANTACY", "THRILLER", "ACTION", "MYTHOLOGY", "ADVENTURE"]
+
+
+def load_metadata(id, name, location, genre):
+    states = (
+        gr.Column(visible=False),
+        gr.Column(visible=True),
+    )
+
+    comic = pipeline.Comic(
+        id=id,
+        name=name,
+        location=location,
+        genre=genre,
+    )
+    return comic, "Generate Comic Synopsis", *states, "synopsis"
+
+
+
+
+character_profiles_format = """\
+### {{name}}<br>
+Age: {{age}}&nbsp;&nbsp;Gender: {{gender}}<br>
+Bio:<br>
+{{description}}<br>
+======================================<br><br>"""
+
+common_format = """\
+### {{text1}}: {{text2}}<br>
+=======================================<br>
+{{synopsis}}<br>
+
+{% if isFrame -%}
+Narration: {{narration}}
+Location: {{location}}
+Frame Setting: {{setting}}
+Dilouge:: {{char}}: '{{dilouge}}'
+{% endif %}<br><br>
+<br>"""
+
+def format_output(output_type: str, output: ..., episode_idx: int = -1, scene_idx: int = -1):
+    if output_type == "synopsis":
+        return textwrap.dedent(f"""\
+## Synopsis:
+{output}""")
+    elif output_type == "episode_breakdown":
+        output_string = ""
+        idxs = sorted([idx for idx, _ in output.items()])
+        for idx in idxs:
+            episode = output[idx]
+            text = jinja2.Template(source=common_format).render({
+                    "text1": idx,
+                    "text2": episode.title,
+                    "synopsis": episode.abstract,        
+                }
+            )
+            output_string += text
+        return output_string
+    elif output_type == "character_generation":
+        output_string = ""
+        for name, profile in output.items():
+            text = jinja2.Template(source=character_profiles_format).render({
+                "name": name,
+                "age": profile.age,
+                "gender": profile.gender,
+                "description": profile.description,
+            })
+            output_string += text
+        return output_string
+    elif output_type == "scenes":
+        output_string = ""
+        for episode_num, episode in output.items():
+            if episode_num != episode_idx:
+                continue
+            for scene_num, scene in enumerate(episode.scenes):
+                text = jinja2.Template(source=common_format).render({
+                    "text1": "Scene",
+                    "text2": scene_num,
+                    "synopsis": scene.description,
+                })
+                output_string += text
+        return output_string
+    elif output_type == "frames":
+        output_string = ""
+        for episode_num, episode in output.items():
+            if episode_num != episode_idx:
+                continue
+            for scene_num, scene in enumerate(episode.scenes):
+                if scene_num != scene_idx:
+                    continue
+                for frame_num, frame in enumerate(scene.frames):
+                    text = jinja2.Template(source=common_format).render(
+                        text1="Frame",
+                        text2=frame_num,
+                        synopsis=frame.description,
+                        isFrame=True,
+                        narration=frame.narration,
+                        location=frame.location,
+                        setting=frame.frame_setting,
+                        char=frame.audio_cue_character,
+                        dilouge=frame.audio_cue_text,
+                    )
+                    output_string += text
+        return output_string
+
+
+def generate_llm_output(
+    user_input: str,
+    system_instruction: str,
+    current_status: str,
+    currComic: pipeline.Comic,
+    episode_idx: int,
+    scene_idx: int,
+    regenerate: bool = False,
+):
+    model_output = ""
+    is_scrollable = False
+    print(f"Pipeline is at: {current_status}")
+    if current_status == "synopsis":
+        currComic.generate_comic_synopsis(user_input, system_instruction, regenerate)
+        model_output = format_output(current_status, currComic.abstract)
+    elif current_status == "episode_breakdown":
+        currComic.generate_episode_breakdowns(system_instruction, regenerate)
+        model_output = format_output(current_status, currComic.episodes)
+    elif current_status == "character_generation":
+        currComic.generate_character_profiles(system_instruction, regenerate)
+        model_output = format_output(current_status, currComic.characters)
+    elif current_status == "scenes":
+        currComic.generate_scenes(
+            system_instruction, episode_idx, regenerate
+        )
+        model_output = format_output(
+            current_status, currComic.episodes, episode_idx
+        )
+        is_scrollable = True
+    elif current_status == "frames":
+        currComic.generate_frames(
+            system_instruction, episode_idx, scene_idx, regenerate
+        )
+        model_output = format_output(
+            current_status, currComic.episodes, episode_idx, scene_idx
+        )
+        is_scrollable = True
+
+    return (
+        gr.Button(visible=is_scrollable),  # Prev Button
+        gr.Button(visible=is_scrollable),  # Next button
+        gr.Button(visible=False),  # Generate Button
+        gr.Button(visible=True),  # Regenerate Button
+        gr.Markdown(value=model_output, visible=True),  # Model output
+        gr.Button(visible=True),  # Submit Button
+    )
+
+
+def submit_data(
+    current_status: str, currComic: pipeline.Comic, episode_idx: int
+):
+    currComic.save_comic_to_s3()
+    preamble = ""
+
+    if current_status == "synopsis":
+        current_status = "episode_breakdown"
+        preamble = script_gen.generate_episode_breakdown_instruction
+    elif current_status == "episode_breakdown":
+        current_status = "character_generation"
+        preamble = script_gen.generate_character_profile_instruction
+    elif current_status == "character_generation":
+        current_status = "scenes"
+        preamble = script_gen.generate_scene_instruction
+        episode_idx = min(list(currComic.episodes.keys()))
+    elif current_status == "scenes":
+        current_status = "frames"
+        preamble = script_gen.generate_frames_instruction
+        episode_idx = min(list(currComic.episodes.keys()))
+    elif current_status == "frames":
+        current_status = "finished"
+    print(f"Next Pipeline stage is: {current_status}")
+    return (
+        episode_idx,
+        current_status,
+        gr.Textbox(visible=False),  # user_input
+        gr.Textbox(value=preamble, visible=True),  # preamble
+        gr.Button(visible=False),  # Prev Button
+        gr.Button(visible=False),  # Next button
+        gr.Button(visible=True),  # Generate Button
+        gr.Button(visible=False),  # Regenerate Button
+        gr.Markdown(value="", visible=False),  # Model output
+    )
+
+
+with gr.Blocks() as demo:
+    currComic = gr.State(None)
+    current_status = gr.State("")
+    episode_idx = gr.State(-1)
+    scene_idx = gr.State(0)
+    regenerate = gr.State(True)
+
+    instruction = gr.Markdown(value="Enter the comic meta data.")
+    with gr.Column(visible=True) as metadata:
+        with gr.Row():
+            comic_id = gr.Textbox(label="Enter Comic ID:", placeholder="Enter Comic ID")
+            comic_name = gr.Textbox(
+                label="Enter Comic Title:", placeholder="Enter Comic Title"
+            )
+            comic_location = gr.Textbox(
+                label="Enter Comic Location:", placeholder="Enter Comic Location"
+            )
+            comic_genre = gr.Dropdown(
+                choices=GENRE_CHOICES, label="Current Genre", interactive=True
+            )
+        metadata_btn = gr.Button("Submit Metadata")
+
+    with gr.Column(visible=False) as synopsis:
+        with gr.Row():
+            user_input = gr.Textbox(label="Enter Your Idea for the comic:")
+            preamble = gr.Textbox(
+                label="Model Instruction:",
+                value=script_gen.generate_synopsis_instruction,
+                lines=10,
+            )
+
+        with gr.Row() as navigation:
+            previous_btn = gr.Button("Previous", visible=False)
+            generate_btn = gr.Button("Generate")
+            regenerate_btn = gr.Button("Re-Generate", visible=False)
+            next_btn = gr.Button("Next", visible=False)
+
+        model_response = gr.Markdown(label="Model Output")
+        submit_btn = gr.Button("Submit", visible=False)
+
+        epi = gr.Textbox(label="Current Episode Synopsis", visible=False)
+        scene = gr.Textbox(label="Current Scene Synopsis", visible=False)
+
+    metadata_btn.click(
+        load_metadata,
+        inputs=[comic_id, comic_name, comic_location, comic_genre],
+        outputs=[currComic, instruction, metadata, synopsis, current_status],
+    )
+
+    generate_btn.click(
+        generate_llm_output,
+        inputs=[
+            user_input,
+            preamble,
+            current_status,
+            currComic,
+            episode_idx,
+            scene_idx
+        ],
+        outputs=[
+            previous_btn,
+            next_btn,
+            generate_btn,
+            regenerate_btn,
+            model_response,
+            submit_btn
+        ],
+    )
+
+    regenerate_btn.click(
+        generate_llm_output,
+        inputs=[
+            user_input,
+            preamble,
+            current_status,
+            currComic,
+            episode_idx,
+            scene_idx,
+            regenerate, # regenrate option
+        ],
+        outputs=[
+            previous_btn,
+            next_btn,
+            generate_btn,
+            regenerate_btn,
+            model_response,
+            submit_btn
+        ],
+    )
+
+    submit_btn.click(
+        submit_data,
+        inputs=[
+            current_status,
+            currComic,
+            episode_idx,
+        ],
+        outputs=[
+            episode_idx,
+            current_status,
+            user_input,
+            preamble,
+            previous_btn,
+            next_btn,
+            generate_btn,
+            regenerate_btn,
+            model_response
+        ]
+    )
+
+demo.launch(auth=("admin", "Qrt@12*34#immersfy"), share=True, ssr_mode=False)

pipeline.py

@@ -0,0 +1,404 @@
+import abc
+import jinja2
+import base64
+import io
+import random
+import dataclasses
+import traceback as tc
+from dacite import from_dict
+from src import datastructures
+from src import aws_utils
+from src import parameters
+from src import script_gen
+from src import grok_wrapper
+
+logger = parameters.LOGGER
+Episode = datastructures.Episode
+Scene = datastructures.Scene
+Frame = datastructures.Frame
+Composition = datastructures.Composition
+Character = datastructures.Character
+Location = datastructures.Location
+llm = grok_wrapper.GROK_2
+
+
+class Comic(abc.ABC):
+
+    def __init__(
+        self,
+        id: str,
+        name: str,
+        location: str,
+        genre: str,
+    ):
+        self.id = id
+        self.name = name
+        self.abstract = ""
+        # self.abstract = abstract
+        self.genre = genre
+        self.location = location
+        self._post_init()
+
+    def _post_init(self):
+        self.episodes = self._get_current_episodes()
+        self.characters = self._get_current_characters()
+
+    def generate_comic_synopsis(
+        self,
+        user_input: str,
+        system_instruction: str,
+        regenerate: bool = False,
+    ):
+        """Class method to generate synopsis for the comic."""
+        if self.abstract and not regenerate:
+            logger.info(
+                f"Synopsis is aleady generated for the comic: {self.id}. Skipping generation!"
+            )
+            return ""
+
+        try:
+            logger.info("Started generating comic Synopsis!")
+            prompt_dict = {
+                "system": system_instruction,
+                "user": user_input,
+            }
+            synopsis = llm.generate_response(prompt_dict)
+
+            # Save the synopsis to s3 bucket at the path formed by:
+            # f"s3://{parameters.AWS_BUCKET}/{self.id}/synopsis.json"
+            aws_utils.save_to_s3(
+                parameters.AWS_BUCKET,
+                self.id,
+                {"synopsis": synopsis},
+                "synopsis.json",
+            )
+            self.abstract = synopsis
+            logger.info("Completed generating comic synopsis!")
+        except Exception:
+            logger.error(
+                f"Faced an exception while generating comic synopsis! See full traceback below"
+            )
+            logger.error(tc.format_exc())
+            return tc.format_exc()
+
+    def _get_current_episodes(self) -> dict:
+        """Fetches the epiodes stored on S3 if present."""
+        try:
+            response = aws_utils.S3_CLIENT.list_objects_v2(
+                Bucket=parameters.AWS_BUCKET,
+                Prefix=f"{self.id}/episodes/",
+                Delimiter="/",
+            ).get("CommonPrefixes", None)
+
+            # In case the URL is incorrect or ill-formatted.
+            if response is None:
+                logger.warning(f"No episodes found for the comic id: {self.id}")
+                return {}
+
+            episodes = {}
+            for prefix in response:
+                folder = prefix["Prefix"]
+                if "episode" in folder:
+                    json_path = f"s3://{parameters.AWS_BUCKET}/{folder}episode.json"
+                    # The sample value of folder: 301/episodes/episode-10
+                    # Thus, split by / and 2nd index gives episode-10
+                    # then, split by - and last index gives the number 10
+                    idx = int(folder.split("/")[2].split("-")[-1])
+                    episodes[idx] = from_dict(
+                        data_class=Episode,
+                        data=eval(
+                            aws_utils.fetch_from_s3(source=json_path).decode("utf-8")
+                        ),
+                    )
+            logger.info(f"Fetched {len(episodes)} episodes for comic id: {self.id}")
+            return episodes
+        except Exception:
+            logger.error(
+                f"Faced an exception while fetching episode breakdowns! See full traceback below:"
+            )
+            logger.error(tc.format_exc())
+            return {}
+
+    def _get_current_characters(self) -> dict:
+        """Fetches the character profiles stored on S3 if present."""
+        try:
+            character_json = (
+                f"s3://{parameters.AWS_BUCKET}/{self.id}/characters/characters.json"
+            )
+            data = eval(aws_utils.fetch_from_s3(source=character_json).decode("utf-8"))
+            chars = {}
+            for name, profile in data.items():
+                # profile["expressions"] = profile.get("expressions", [])
+                chars[name] = Character(**profile)
+            logger.info(
+                f"Fetched {len(chars)} character profiles for comic id: {self.id}"
+            )
+            return chars
+        except Exception:
+            logger.error(
+                f"Faced an exception while fetching character profiles! See full traceback below:"
+            )
+            logger.error(tc.format_exc())
+            return {}
+
+    def generate_episode_breakdowns(self, system_instruction: str, regenerate: bool = False) -> str:
+        """Class method to generate episode breakdowns for the comic."""
+        if len(self.episodes) > 0 and not regenerate:
+            logger.info(
+                f"Episodes are aleady generated for the comic: {self.id}. Skipping generation!"
+            )
+            return ""
+
+        try:
+            logger.info("Started generating episode breakdowns!")
+            # We employ a 2-step process to generate the abstracts for each episodes.
+            # 1. Initially we take the comic abstract and expand it to include some more creative
+            #    twists and turns for the story.
+            # 2. Next, we use that expanded synopsis to generate per episode abstract
+            #    ensuring the flow of the story doesn't break and each episode is interesting enough
+            #    to read.
+            prompt_dict = {
+                "system": script_gen.generate_detailed_synopsis_instruction,
+                "user": jinja2.Template(
+                    source=script_gen.generate_detailed_synopsis_user_prompt
+                ).render(
+                    {
+                        "TITLE": self.name,
+                        "GENRE": self.genre,
+                        "SYNOPSIS": self.abstract,
+                        "LOCATION": self.location,
+                    }
+                ),
+            }
+            extended_synopsis = llm.generate_response(prompt_dict)
+
+            prompt_dict = {
+                "system": system_instruction,
+                "user": jinja2.Template(
+                    source=script_gen.generate_episode_breakdown_user_prompt
+                ).render({"SYNOPSIS": extended_synopsis}),
+            }
+            episode_breakdown = llm.generate_valid_json_response(prompt_dict)
+            logger.info(
+                f"Generated a total of {len(episode_breakdown)} episodes for comic: {self.id}"
+            )
+
+            # Save the breakdown to s3 bucket at the path formed by:
+            # f"s3://{parameters.AWS_BUCKET}/{self.id}/episode_breakdown.json"
+            episode_breakdown_path = aws_utils.save_to_s3(
+                parameters.AWS_BUCKET,
+                self.id,
+                episode_breakdown,
+                "episode_breakdown.json",
+            )
+
+            for idx, data in episode_breakdown.items():
+                self.episodes[idx] = Episode(
+                    title=data["title"],
+                    abstract=data["synopsis"],
+                )
+            self.save_comic_to_s3()
+
+            logger.info("Completed generating episode breakdowns!")
+            return episode_breakdown_path
+        except Exception:
+            logger.error(
+                f"Faced an exception while generating episode breakdowns! See full traceback below"
+            )
+            logger.error(tc.format_exc())
+            return tc.format_exc()
+
+    def generate_character_profiles(self, system_instruction: str, regenerate: bool = False) -> str:
+        """Class method to generate character profiles for the comic."""
+        if len(self.characters) > 0 and not regenerate:
+            logger.info(
+                f"Characters are aleady generated for the comic: {self.id}. Skipping generation!"
+            )
+            return ""
+
+        try:
+            logger.info("Started generating character profiles!")
+            prompt_dict = {
+                "system": system_instruction,
+                "user": jinja2.Template(
+                    source=script_gen.generate_character_profile_user_prompt
+                ).render(
+                    {
+                        "COMIC_NAME": self.name,
+                        "ABSTRACT": self.abstract,
+                        "episode_breakdown": [
+                            dataclasses.asdict(episode)
+                            for _, episode in self.episodes.items()
+                        ],
+                    }
+                ),
+            }
+            characters = llm.generate_valid_json_response(prompt_dict)
+
+            for char in characters:
+                char["profile_image"] = (
+                    f"s3://blix-demo-v0/{self.id}/characters/images/{char['name']}.jpg"
+                )
+                self.characters[char["name"]] = Character(**char)
+
+            character_profile_path = aws_utils.save_to_s3(
+                parameters.AWS_BUCKET,
+                f"{self.id}/characters",
+                {
+                    name: dataclasses.asdict(char)
+                    for name, char in self.characters.items()
+                },
+                "characters.json",
+            )
+            logger.info(
+                f"Model extracted a total of {len(self.characters)} characters."
+            )
+            logger.info("Completed generating character profiles!")
+            return character_profile_path
+        except Exception:
+            logger.error(
+                f"Faced an exception while generating character profiles! See full traceback below"
+            )
+            logger.error(tc.format_exc())
+            return tc.format_exc()
+
+    def generate_scenes(self, system_instruction: str, episode_idx: int, regenerate: bool = False) -> str:
+        """Class method to generate scenes for the comic."""
+        try:
+            logger.info("Started generating Scenes!")
+            # The idea is to take each of the episodes and break them down into different
+            # scenes while maintaining the overall story flow and character details.
+            for episode_num, episode in self.episodes.items():
+                if len(episode.scenes) > 0 and not regenerate:
+                    logger.info(
+                        f"Scenes are aleady generated for the episode: {episode_num}. Skipping generation!"
+                    )
+                    continue
+
+                # In case regeneration is enabled, we check if we're at the correct episode index.
+                if regenerate and episode_num != episode_idx:
+                    continue
+                logger.info(f"Generating Scenes for episode: {episode_num}")
+                prompt_dict = {
+                    "system": system_instruction,
+                    "user": jinja2.Template(
+                        source=script_gen.generate_scene_user_prompt
+                    ).render(
+                        {
+                            "SYNOPSIS": episode.abstract,
+                            "CHARS": [
+                                dataclasses.asdict(char)
+                                for _, char in self.characters.items()
+                            ],
+                        }
+                    ),
+                }
+                episode_scenes = llm.generate_valid_json_response(prompt_dict)
+
+                episode.scenes = [
+                    Scene(description=description)
+                    for _, description in episode_scenes["scenes"].items()
+                ]
+                _ = aws_utils.save_to_s3(
+                    parameters.AWS_BUCKET,
+                    f"{self.id}/episodes/episode-{episode_num}",
+                    episode_scenes,
+                    "scenes.json",
+                )
+            self.save_comic_to_s3()
+            scenes_path = (
+                f"s3://{parameters.AWS_BUCKET}/{self.id}/episodes/episode-*/scenes.json"
+            )
+            logger.info("Completed generating Scenes!")
+            return scenes_path
+        except Exception:
+            logger.error(
+                f"Faced an exception while generating scenes! See full traceback below"
+            )
+            logger.error(tc.format_exc())
+            return tc.format_exc()
+
+    def generate_frames(self, system_instruction: str, episode_idx: int, scene_idx: int, regenerate: bool = False) -> str:
+        """Class method to generate frames for the comic."""
+        try:
+            location_list = []
+            logger.info("Started generating frames!")
+            # Next step is to take each of the generated scenes and break them down
+            # further into various frames. We target to achieve a specific still moment
+            # in the comic book story that resembles a frame and then try to visualize it.
+            for episode_num, episode in self.episodes.items():
+                for scene_num, scene in enumerate(episode.scenes):
+                    if len(scene.frames) > 0 and not regenerate:
+                        logger.info(
+                            f"Frames are aleady generated for the episode: {episode_num} and scene: {scene_num}. Skipping generation!"
+                        )
+                        continue
+                
+                    # In case regeneration is enabled, we check if we're at the correct (episode, scene) index.
+                    if regenerate and (episode_num != episode_idx or scene_num != scene_idx):
+                        continue
+                    logger.info(
+                        f"Generating frames for episode: {episode_num} and scene: {scene_num}"
+                    )
+                    prompt_dict = {
+                        "system": system_instruction,
+                        "user": jinja2.Template(
+                            source=script_gen.generate_frames_user_prompt
+                        ).render(
+                            {
+                                "SYNOPSIS": scene.description,
+                                "CHARACTERS": [
+                                    dataclasses.asdict(char)
+                                    for _, char in self.characters.items()
+                                ],
+                            }
+                        ),
+                    }
+                    frames = llm.generate_valid_json_response(prompt_dict)
+                    for frame in frames:
+                        chars = []
+                        for char in frame["characters"]:
+                            if self.characters.get(char, None) is not None:
+                                chars.append(self.characters[char])
+                        frame["characters"] = chars
+
+                    scene.frames = [Frame(**frame) for frame in frames["images"]]
+                    aws_utils.save_to_s3(
+                        parameters.AWS_BUCKET,
+                        f"{self.id}/episodes/episode-{episode_num}",
+                        {
+                            "frames": [
+                                dataclasses.asdict(frame) for frame in scene.frames
+                            ]
+                        },
+                        f"scene-{scene_num}-frames.json",
+                    )
+            self.save_comic_to_s3()
+            frames_path = f"s3://{parameters.AWS_BUCKET}/{self.id}/episodes/episode-*/scene-*-frames.json"
+            logger.info("Completed generating frames!")
+            return frames_path
+        except Exception:
+            logger.error(
+                f"Faced an exception while generating frames! See full traceback below"
+            )
+            logger.error(tc.format_exc())
+            return tc.format_exc()
+
+    def save_comic_to_s3(self) -> bool:
+        """Class method to save data to s3 bucket."""
+        try:
+            for episode_num, episode in self.episodes.items():
+                logger.info(f"Saving episode: {episode_num}")
+                path = aws_utils.save_to_s3(
+                    parameters.AWS_BUCKET,
+                    f"{self.id}/episodes/episode-{episode_num}",
+                    dataclasses.asdict(episode),
+                    "episode.json",
+                )
+            return True
+        except Exception:
+            logger.error(
+                f"Faced an exception while saving comic to S3! See full traceback below"
+            )
+            logger.error(tc.format_exc())
+            return False

src/aws_utils.py

@@ -0,0 +1,136 @@
+import os
+import json
+from io import BytesIO
+from typing import Union
+from urllib.parse import urlparse
+
+import boto3
+from botocore.client import Config
+from botocore.exceptions import NoCredentialsError
+
+AWS_REGION = os.getenv("AWS_REGION")
+
+# Initialize the S3 client
+S3_CLIENT = boto3.client(
+    "s3", region_name=AWS_REGION, config=Config(signature_version="s3v4")
+)
+
+
+def save_to_s3(
+    bucket_name: str,
+    folder_name: str,
+    content: Union[str, dict, BytesIO],
+    file_name: str,
+) -> str:
+    """
+    Save a file to an S3 bucket, determining the content type based on the input type.
+
+    Args:
+        bucket_name (str): The name of the S3 bucket.
+        folder_name (str): The folder path in the S3 bucket.
+        content (Union[str, dict, BytesIO]): The content to save, can be a string, dictionary, or BytesIO.
+        file_name (str): The file name under which the content should be saved.
+
+    Returns:
+        str: The S3 URL of the uploaded file, or an error message if credentials are not available.
+    """
+    # Ensure the folder name ends with a '/'
+    # if not folder_name.endswith('/'):
+    #     folder_name += '/'
+    # Determine file name and content type based on the input
+    if isinstance(content, str):
+        file_content = content
+        content_type = "text/plain"
+    elif isinstance(content, dict):
+        file_content = json.dumps(content)
+        content_type = "application/json"
+    elif isinstance(content, BytesIO):
+        file_content = content
+        content_type = "image/jpeg"
+    else:
+        print(
+            "Invalid content type. Content must be a string, dictionary, or BytesIO."
+        )
+        raise ValueError("Content must be either a string, dictionary, or BytesIO.")
+
+    # Ensure the folder name ends with a '/'
+    s3_file_path = f"{folder_name.rstrip('/')}/{file_name}"
+
+    try:
+        # Upload the file to S3
+        S3_CLIENT.put_object(
+            Bucket=bucket_name,
+            Key=s3_file_path,
+            Body=file_content,
+            ContentType=content_type,
+        )
+        s3_url = f"s3://{bucket_name}/{s3_file_path}"
+        print(f"File successfully uploaded to {s3_url}")
+        return s3_url
+
+    except NoCredentialsError:
+        print("AWS credentials not available.")
+        return "Error: AWS credentials not available."
+
+
+def fetch_from_s3(source: Union[str, dict], region_name: str = "ap-south-1") -> bytes:
+    """
+    Fetch a file's content from S3 given a source URL or dictionary with bucket and key.
+
+    Args:
+        source (Union[str, dict]): The source S3 URL or a dictionary with 'bucket_name' and 'file_key'.
+        region_name (str): The AWS region name for the S3 client (default is 'ap-south-1').
+
+    Returns:
+        bytes: The content of the file fetched from S3.
+    """
+    print(f"Fetching file from S3. Source: {source}")
+    s3_client = boto3.client("s3", region_name=region_name)
+
+    # Parse the source depending on its type
+    if isinstance(source, str):
+        parsed_url = urlparse(source)
+        bucket_name = parsed_url.netloc.split(".")[0]
+        file_path = parsed_url.path.lstrip("/")
+    elif isinstance(source, dict):
+        bucket_name = source.get("bucket_name")
+        file_path = source.get("file_key")
+        if not bucket_name or not file_path:
+            print("Dictionary input must contain 'bucket_name' and 'file_key'.")
+            raise ValueError(
+                "Dictionary input must contain 'bucket_name' and 'file_key'."
+            )
+    else:
+        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}")
+    try:
+        response = s3_client.get_object(Bucket=bucket_name, Key=file_path)
+        file_content = response["Body"].read()
+        print(f"File fetched successfully from {bucket_name}/{file_path}")
+        return file_content
+    except Exception as e:
+        print(f"Failed to fetch file from S3: {e}")
+        raise
+
+
+def list_s3_objects(bucket_name: str, folder_path: str = "") -> list:
+    """
+    Lists a content of the given a directory URL.
+
+    Args:
+        bucket_name (str): The name of the S3 bucket.
+        folder_name (str): The folder path in the S3 bucket.
+
+    Returns:
+        list: The list of files found inside the given directory URL.
+    """
+    response = S3_CLIENT.list_objects_v2(Bucket=bucket_name, Prefix=folder_path)
+    # Check if the bucket contains objects
+    objects = []
+    if "Contents" in response:
+        for obj in response["Contents"]:
+            objects.append(obj["Key"])
+
+    return objects

src/custom_log.py

@@ -0,0 +1,52 @@
+import logging
+import sys
+
+import colorlog
+
+
+def setup_logger(name: str) -> logging.Logger:
+    """
+    Set up a logger with colored output for console if running in a terminal.
+
+    Args:
+        name (str): Name of the logger.
+
+    Returns:
+        logging.Logger: Configured logger instance.
+    """
+    # Check if the output is a terminal to enable colored output
+    use_colors = sys.stdout.isatty()
+
+    # Create a stream handler
+    handler = colorlog.StreamHandler()
+
+    if use_colors:
+        # Use colored formatter for terminal output
+        handler.setFormatter(
+            colorlog.ColoredFormatter(
+                "%(asctime)s - %(log_color)s%(levelname)s%(reset)s - %(message)s",
+                datefmt="%Y-%m-%d %H:%M:%S",
+                log_colors={
+                    "DEBUG": "cyan",
+                    "INFO": "green",
+                    "WARNING": "yellow",
+                    "ERROR": "red",
+                    "CRITICAL": "bold_red",
+                },
+            )
+        )
+    else:
+        # Use standard formatter when not in a terminal
+        handler.setFormatter(
+            logging.Formatter(
+                "%(asctime)s - %(levelname)s - %(message)s", datefmt="%Y-%m-%d %H:%M:%S"
+            )
+        )
+
+    # Configure the logger
+    logger = logging.getLogger(name)
+    logger.addHandler(handler)
+    logger.setLevel(logging.INFO)
+    logger.propagate = False
+
+    return logger

src/datastructures.py

@@ -0,0 +1,55 @@
+from typing import List, Dict
+import dataclasses
+
+
+@dataclasses.dataclass
+class Composition:
+    prompt: str
+    shot_type: str
+    seed: int = None
+    image: str = ""
+
+
+@dataclasses.dataclass
+class Character:
+    name: str
+    age: str
+    gender: str
+    description: str
+    profile_image: str
+    expressions : List[str] = dataclasses.field(default_factory=list)
+    compositions: List[Composition] = dataclasses.field(default_factory=list)
+
+
+@dataclasses.dataclass(frozen=True)
+class Location:
+    name: str
+    description: str
+
+@dataclasses.dataclass
+class Frame:
+    description: str
+    narration: str
+    audio_cue_text: str
+    audio_cue_character: str
+    location: str
+    frame_setting: str
+    characters: List[Character] = dataclasses.field(default_factory=list)
+    compositions: List[Composition] = dataclasses.field(default_factory=list)
+    character_expression: Dict[str, str] = dataclasses.field(default_factory=dict)
+    audio_path: str = ""
+    stitch_image: str = ""
+
+@dataclasses.dataclass
+class Scene:
+    description: str
+    frames: List[Frame] = dataclasses.field(default_factory=list)
+
+
+@dataclasses.dataclass
+class Episode:
+    title: str
+    abstract: str
+    scenes: List[Scene] = dataclasses.field(default_factory=list)
+    thumbnail_image: str  = ""
+    thumbnail_image_prompt: str = ""

src/grok_wrapper.py

@@ -0,0 +1,81 @@
+"""Model wrapper to interact with OpenAI models."""
+import abc
+import ast
+from typing import Mapping
+
+import openai
+
+from src import parameters
+
+logger = parameters.LOGGER
+
+
+class xAIModel(abc.ABC):
+    API_KEY = ""
+
+    def __init__(self, model_name: str, API_KEY: str):
+        try:
+            self.client = openai.OpenAI(
+                api_key=API_KEY,
+                base_url="https://api.x.ai/v1",
+            )
+            self.model_name = model_name
+        except Exception as exc:
+            raise Exception(
+                "Failed to initialize Grok xAI model client. See traceback for more details.",
+            ) from exc
+
+    def prepare_input(self, prompt_dict: Mapping[str, str]) -> str:
+        conversation = []
+        try:
+            for role, content in prompt_dict.items():
+                conversation.append({"role": role, "content": content})
+            return conversation
+        except Exception as exc:
+            raise Exception(
+                f"Incomplete Prompt Dictionary Passed. Expected to have atleast a role and it's content.\nPassed dict: {prompt_dict}",
+            ) from exc
+
+    def generate_response(
+        self,
+        prompt_dict: Mapping[str, str],
+        max_output_tokens: int = None,
+        temperature: int = 0.6,
+        response_format: dict = None,
+    ) -> str:
+        conversation = self.prepare_input(prompt_dict)
+        try:
+            response = self.client.chat.completions.create(
+                model=self.model_name,
+                messages=conversation,
+                max_tokens=max_output_tokens if max_output_tokens else None,
+                temperature=temperature,
+                response_format=response_format,
+            )
+            return response.choices[0].message.content
+        except Exception as exc:
+            raise Exception(
+                f"Exception in generating model response.\nModel name: {self.model_name}\nInput prompt: {str(conversation)}",
+            ) from exc
+
+    def generate_valid_json_response(
+        self,
+        prompt_dict: Mapping[str, str],
+        max_output_tokens: int = None,
+        temperature: int = 0.6,
+    ) -> str:
+        """Generate a response with retries, returning a valid JSON."""
+        for _ in range(parameters.MAX_TRIES):
+            try:
+                model_response = self.generate_response(
+                    prompt_dict, max_output_tokens, temperature, {"type": "json_object"}
+                )
+                return ast.literal_eval(model_response)
+            except Exception as e:
+                continue
+        raise Exception(
+            f"Maximum retries met before valid JSON structure was found.\nModel name: {self.model_name}\nInput prompt: {str(prompt_dict)}"
+        )
+
+
+GROK_2 = xAIModel("grok-2-1212", parameters.XAI_API_KEY)

src/parameters.py

@@ -0,0 +1,24 @@
+import os
+
+from dotenv import load_dotenv
+
+from src import custom_log
+
+load_dotenv()
+
+# AWS Resouce variables
+os.environ["AWS_ACCESS_KEY_ID"] = os.environ.get("AWS_ACCESS_KEY_ID")
+os.environ["AWS_SECRET_ACCESS_KEY"] = os.environ.get("AWS_SECRET_ACCESS_KEY")
+os.environ["S3_BUCKET_NAME"] = os.environ.get("AWS_BUCKET")
+AWS_BUCKET = os.environ.get("AWS_BUCKET")
+AWS_REGION = os.environ.get("AWS_REGION")
+
+# Local variables
+MAX_WORKERS = int(os.environ.get("MAX_WORKERS"))
+MAX_TRIES = int(os.environ.get("MAX_TRIES"))
+
+# xAI AI API Key
+XAI_API_KEY = os.environ.get("XAI_API_KEY")
+
+# Custom logger
+LOGGER = custom_log.setup_logger(__name__)

src/script_gen.py

@@ -0,0 +1,204 @@
+"""Collection of all system prompts for `/generate_episode_breakdown` and `/generate_episode_content` API."""
+
+###################################### STEP 0: GENERATE COMIC SYNOPSIS ######################################
+
+generate_synopsis_instruction = """\
+You're an imaginative comic book creator with a passion for storytelling and a keen eye for detail. You have been crafting engaging comic synopses and plots for over a decade, ensuring each story is captivating and resonates with readers of all ages.
+
+You're given details about a comic book idea along with it's genre, title and location. Your task is to generate a synopsis for it.
+Follow these guidelines:
+1) Ensure that the synopsis incorporates interesting character arcs, vibrant settings, and engaging conflicts that make the story exciting and fun to read.
+2) Only output the synopsis as a single text field.
+3) Make sure to include hints about the given location in the story for ex if the Story location is Mumbai, include any of the iconic places like Gateway of India, Nariman point, Marine drive, etc.
+"""
+
+###################################### STEP 1: GENERATE EPISODE BREAKDOWNS ######################################
+
+# System Instruction to extend the given comic synopsis.
+generate_detailed_synopsis_instruction = """\
+You're an imaginative comic book creator with a passion for storytelling and a keen eye for detail. You have been crafting engaging comic synopses and plots for over a decade, ensuring each story is captivating and resonates with readers of all ages.
+
+You're given details about a comic book idea like it's short synopsis, genre, title and location. Your task is to expand and enhance it further.
+Follow these guidelines:
+1) Ensure that the synopsis incorporates interesting character arcs, vibrant settings, and engaging conflicts that make the story exciting and fun to read.
+2) Only output the synopsis as a single text field.
+3) Make sure to include hints about the given location in the story for ex if the Story location is Mumbai, include any of the iconic places like Gateway of India, Nariman point, Marine drive, etc.
+"""
+
+generate_detailed_synopsis_user_prompt = """\
+Here are the details for the comic:
+- Title: {{TITLE}}
+- Genre: {{GENRE}}
+- Synopsis: {{SYNOPSIS}}
+- Location: {{LOCATION}}
+"""
+
+# System Instruction to generate episode breakdown from Comic Abstract, Genre etc
+generate_episode_breakdown_instruction = """\
+You're a seasoned comic book writer and plot developer with extensive experience in crafting engaging and complex comic episode breakdowns that captivate readers.
+You have a keen ability to weave intricate character relationships, compelling subplots, and emotional arcs, all while ensuring that each episode contributes meaningfully to the overall narrative.
+
+You're given a detailed synopsis about a comic book along with the location and your task is to generate a per episode breakdown of how the story unfolds.
+Follow these guidelines:
+1) The breakdown should encompass the evolution of character relationships, include supporting or temporary characters, and incorporate subplots that enhance the reading experience.
+2) Ensure proper flow of the story across the episodes i.e. each new episode's start should be coherent with the last episode's ending.
+3) Pay close attention to cohesion with the overall plot, character consistency, thematic continuity, adaptability, and reader engagement.
+4) Make sure to incorporate hints about the location in the story for example if the story location is New York include any places like Manhattan bridge, Empire state building, etc.
+4) Generate the breakdown in a JSON format with the below format:
+{
+    "1": {
+        "title" : "title of episode 1",
+        "synopsis" : "synopsis of episode 1",
+    },
+    "2":{
+        "title" : "title of episode 2",
+        "synopsis" : "synopsis of episode 2",
+    },
+    ...
+}
+5) The total number of episodes should be in the range of 5-15.
+
+Remember that the JSON output must be valid and should not contain any special characters, single quotes, backslashes, or any additional text that could invalidate the JSON format.
+"""
+
+# User Prompt to generate episode breakdown from Comic Abstract, Genre etc
+generate_episode_breakdown_user_prompt = (
+    """Here is the synopsis of the comic: {{SYNOPSIS}}"""
+)
+
+
+###################################### STEP 2: GENERATE ALL CHARACTERS ######################################
+
+generate_character_profile_instruction = """\
+You are a creative comic artist and writer with a knack for creative writing and direction.
+
+You will be given details about a comic along with an episode breakdown in a JSON format:
+{   # Each key, value pair inside this dictionary contains a description of each episode's storyline.
+    "1": {
+        "title" : "title of episode 1",
+        "abstract" : "synopsis of episode 1",
+        ... other details
+    },
+    "2":{
+        "title" : "title of episode 2",
+        "abstract" : "synopsis of episode 2",
+        ... other details
+    },
+    ...
+}
+
+Your task is to analyse this JSON and once you've interpret it, then you have to output character profiles for all the characters present in the episodes mentioned.
+Here's a step-by-step plan for you to proceed with some guidelines:
+1) Analyse the JSON and understand the story.
+2) Identify all the different characters appearing in the given episode descriptions.
+3) Make sure that each character dictionary should belong to ony ONE character and not multiple characters like 'children' or 'Ancient Guardians', etc. Give each and every character a separate dictionary.
+4) For every single character prepare a profile in the below format:
+{
+    "name": "Name of the character",
+    "age": "Age of the character",
+    "gender":"Gender of the character",
+    "description":"A 1-liner description describing character's physical apperance like their eyes, face features, height, how they dress, etc."
+}
+4) Output a list of such Dictionaries denoting all the characters found in the whole comic.
+5) Each character should have a unique and natural name which fits with the story line. Do not create characters profiles like "The Group", "Group Member 1", etc.
+6) In case of non human characters, give the profile a proper name suited to it's characteristics and give major focus in create a correct and apt desription.
+7) Avoid creating profiles for temporary characters and inanimate objects whose influence is very less in the story line.
+Keep in mind that the JSON output must be valid and free of any special characters, single quotes, backslashes, or other text that could compromise the format.
+"""
+
+generate_character_profile_user_prompt = """\
+The comic details are as below:
+Name: {{COMIC_NAME}}
+Abstract: {{ABSTRACT}}
+
+Episode breakdown:
+{{episode_breakdown}}"""
+
+
+###################################### STEP 3: GENERATE ALL SCENES ######################################
+
+generate_scene_instruction = """\
+You are an experienced comic artist and narrative designer with deep expertise in sequential storytelling and dramatic writing. Your task is to create rich, narrative scene descriptions for comic episodes that read like engaging story segments while maintaining continuity and emotional depth.
+
+You have been given:
+1. Current episode synopsis (to be broken into scenes)
+2. Character information
+
+Your approach should include:
+1. Analyze the episode synopsis to identify key story beats and emotional arcs
+2. Understand character relationships and their dynamics
+3. Break down the episode into key narrative moments
+4. For each scene, craft a compelling narrative description that:
+   - Flows naturally like a story segment
+   - Weaves setting details into the action
+   - Reveals character emotions through their actions and interactions
+   - Captures dialog moments and character dynamics
+   - Creates proper pacing and builds tension organically
+   - Maintains story momentum while allowing for artistic interpretation
+   - Create clear transitions between scenes that feel story-driven
+5. Generate a total of 3-7 scenes for each episode breakdown.
+
+Output Format in below json:
+{
+   "scenes": {
+       "Scene_1": "A story-like narrative description flowing naturally through the key story beats, character interactions, and emotional moments...",
+       "Scene_2": "Continuation of the story through the next pivotal moment...",
+       "Scene_3": "Further story development with natural progression...",
+       ...
+   }
+}
+"""
+
+generate_scene_user_prompt = """\
+Here's the current episode synopsis:
+
+## Synopsis:
+{{SYNOPSIS}}
+
+## Characters:
+{{CHARS}}
+"""
+
+
+###################################### STEP 4: GENERATE ALL FRAMES ######################################
+
+generate_frames_instruction = """\
+Your task is to transform a given scene synopsis into a series of distinct, sequential image frames that capture the key moments of the story while maintaining continuity.
+Follow the rules below:
+
+1) Output JSON with the following structure for each image:
+{
+    "images": [
+        {
+            "description": "detail visual description of the frame",
+            "narration": "Key story text representing the current frame flow",
+            "characters": ["List of characters in the frame"],
+            "audio_cue_text": "Concise dialogue (max 20 words)",
+            "audio_cue_character": "Character name which speaks the dilouge",
+            "location": "Description about the background where the specific frame is taking place.",
+            "frame_setting": "Description about the current setting of the frame via features like atmosphere, lighting, is it day-light or at night, if the wind is blowing, if the sun is shining, if it is raining, etc. Add elements which help a director film the exact frame at the given location.",
+            "character_expression" : {"character name" : "Expression","character name2": "Expression"} expressions must be from (CRYING, HAPPY, OVERJOYED, CALM, FEAR)
+        }
+        ...
+    ]
+}
+2) Ensure the image frames are interesting for readers.
+3) Each image description must depict a specific still moment in the story. Do not combine multiple actions into a single image.
+4) Establish continuity between images. For example, if the last scene shows Arjun holding his helmet, the next might depict "Arjun's helmet lying on the street," creating an engaging connection that implies Arjun dropped it.
+5) Be creative, but do not fabricate details that conflict with the given storyline.
+6) Aim to generate 10-20 image descriptions based on the source material.
+7) When generating dilouges make them conversational i.e. the dilouge of frame 1 should coincide with dilouge of frame 2 and so on.
+8) In the frame_setting key, do not mention aything about the location, just mention features about the frame which help establish it like atmosphere, lighting, is it day-light or at night, if the wind is blowing, if the sun is shining, if it is raining, etc.
+8) When generating location make sure to include a very specific place where the current frame is taking place, for ex: In middle of New York Time Square.
+9) Make sure that when more than 1 frames are being filmed at the same location, use the same location key text to address them. Use frame_setting key to highlight differences in in-animate conditions across those frames.
+"""
+
+generate_frames_user_prompt = """\
+Here's the current scene synopsis:
+
+## Synopsis:
+{{SYNOPSIS}}
+
+##Characters:
+{{CHARACTERS}}
+"""