malepati/custom_template_working

1import asyncio2from typing import List, Tuple3from models.image_prompt import ImagePrompt4from models.sql.image_asset import ImageAsset5from models.sql.slide import SlideModel6from services.icon_finder_service import ICON_FINDER_SERVICE7from services.image_generation_service import ImageGenerationService8from utils.asset_directory_utils import get_images_directory9from utils.dict_utils import get_dict_at_path, get_dict_paths_with_key, set_dict_at_path10 11 12async def process_slide_and_fetch_assets(13    image_generation_service: ImageGenerationService,14    slide: SlideModel,15) -> List[ImageAsset]:16 17    async_tasks = []18 19    image_paths = get_dict_paths_with_key(slide.content, "__image_prompt__")20    icon_paths = get_dict_paths_with_key(slide.content, "__icon_query__")21 22    for image_path in image_paths:23        __image_prompt__parent = get_dict_at_path(slide.content, image_path)24        async_tasks.append(25            image_generation_service.generate_image(26                ImagePrompt(27                    prompt=__image_prompt__parent["__image_prompt__"],28                )29            )30        )31 32    for icon_path in icon_paths:33        __icon_query__parent = get_dict_at_path(slide.content, icon_path)34        async_tasks.append(35            ICON_FINDER_SERVICE.search_icons(__icon_query__parent["__icon_query__"])36        )37 38    results = await asyncio.gather(*async_tasks)39    results.reverse()40 41    return_assets = []42    for image_path in image_paths:43        image_dict = get_dict_at_path(slide.content, image_path)44        result = results.pop()45        if isinstance(result, ImageAsset):46            return_assets.append(result)47            image_dict["__image_url__"] = result.path48        else:49            image_dict["__image_url__"] = result50        set_dict_at_path(slide.content, image_path, image_dict)51 52    for icon_path in icon_paths:53        icon_dict = get_dict_at_path(slide.content, icon_path)54        icon_dict["__icon_url__"] = results.pop()[0]55        set_dict_at_path(slide.content, icon_path, icon_dict)56 57    return return_assets58 59 60async def process_old_and_new_slides_and_fetch_assets(61    image_generation_service: ImageGenerationService,62    old_slide_content: dict,63    new_slide_content: dict,64) -> List[ImageAsset]:65    # Finds all old images66    old_image_dict_paths = get_dict_paths_with_key(67        old_slide_content, "__image_prompt__"68    )69    old_image_dicts = [70        get_dict_at_path(old_slide_content, path) for path in old_image_dict_paths71    ]72    old_image_prompts = [73        old_image_dict["__image_prompt__"] for old_image_dict in old_image_dicts74    ]75 76    # Finds all old icons77    old_icon_dict_paths = get_dict_paths_with_key(old_slide_content, "__icon_query__")78    old_icon_dicts = [79        get_dict_at_path(old_slide_content, path) for path in old_icon_dict_paths80    ]81    old_icon_queries = [82        old_icon_dict["__icon_query__"] for old_icon_dict in old_icon_dicts83    ]84 85    # Finds all new images86    new_image_dict_paths = get_dict_paths_with_key(87        new_slide_content, "__image_prompt__"88    )89    new_image_dicts = [90        get_dict_at_path(new_slide_content, path) for path in new_image_dict_paths91    ]92 93    # Finds all new icons94    new_icon_dict_paths = get_dict_paths_with_key(new_slide_content, "__icon_query__")95    new_icon_dicts = [96        get_dict_at_path(new_slide_content, path) for path in new_icon_dict_paths97    ]98 99    # Creates async tasks for fetching new images100    async_image_fetch_tasks = []101    new_images_fetch_status = []102 103    # Creates async tasks for fetching new icons104    async_icon_fetch_tasks = []105    new_icons_fetch_status = []106 107    # Creates async tasks for fetching new images108    # Use old image url if prompt is same109    for new_image in new_image_dicts:110        if new_image["__image_prompt__"] in old_image_prompts:111            old_image_url = old_image_dicts[112                old_image_prompts.index(new_image["__image_prompt__"])113            ]["__image_url__"]114            new_image["__image_url__"] = old_image_url115            new_images_fetch_status.append(False)116            continue117 118        async_image_fetch_tasks.append(119            image_generation_service.generate_image(120                ImagePrompt(121                    prompt=new_image["__image_prompt__"],122                )123            )124        )125        new_images_fetch_status.append(True)126 127    # Creates async tasks for fetching new icons128    # Use old icon url if query is same129    for new_icon in new_icon_dicts:130        if new_icon["__icon_query__"] in old_icon_queries:131            old_icon_url = old_icon_dicts[132                old_icon_queries.index(new_icon["__icon_query__"])133            ]["__icon_url__"]134            new_icon["__icon_url__"] = old_icon_url135            new_icons_fetch_status.append(False)136            continue137 138        async_icon_fetch_tasks.append(139            ICON_FINDER_SERVICE.search_icons(new_icon["__icon_query__"])140        )141        new_icons_fetch_status.append(True)142 143    new_images = await asyncio.gather(*async_image_fetch_tasks)144    new_icons = await asyncio.gather(*async_icon_fetch_tasks)145 146    # list of new assets147    new_assets = []148 149    # Sets new image and icon urls for assets that were fetched150    for i, new_image in enumerate(new_images):151        if new_images_fetch_status[i]:152            fetched_image = new_images[i]153            if isinstance(fetched_image, ImageAsset):154                new_assets.append(fetched_image)155                image_url = fetched_image.path156            else:157                image_url = fetched_image158            new_image_dicts[i]["__image_url__"] = image_url159 160    for i, new_icon in enumerate(new_icons):161        if new_icons_fetch_status[i]:162            new_icon_dicts[i]["__icon_url__"] = new_icons[i][0]163 164    for i, new_image_dict in enumerate(new_image_dicts):165        set_dict_at_path(new_slide_content, new_image_dict_paths[i], new_image_dict)166 167    for i, new_icon_dict in enumerate(new_icon_dicts):168        set_dict_at_path(new_slide_content, new_icon_dict_paths[i], new_icon_dict)169 170    return new_assets171 172 173def process_slide_add_placeholder_assets(slide: SlideModel):174 175    image_paths = get_dict_paths_with_key(slide.content, "__image_prompt__")176    icon_paths = get_dict_paths_with_key(slide.content, "__icon_query__")177 178    for image_path in image_paths:179        image_dict = get_dict_at_path(slide.content, image_path)180        image_dict["__image_url__"] = "/static/images/placeholder.jpg"181        set_dict_at_path(slide.content, image_path, image_dict)182 183    for icon_path in icon_paths:184        icon_dict = get_dict_at_path(slide.content, icon_path)185        icon_dict["__icon_url__"] = "/static/icons/placeholder.svg"186        set_dict_at_path(slide.content, icon_path, icon_dict)187