From 0a887bf7a1fbb8f89021af28fb14f5dabdeb961f Mon Sep 17 00:00:00 2001
From: vikramxD <singh.vikram.1782000@gmail.com>
Date: Wed, 25 Dec 2024 18:53:53 +0000
Subject: [PATCH] Add Allegro inference configuration and video generation
 scripts

- Introduced `allegro_settings.py` for managing Allegro model settings with Pydantic.
- Created `allegro_diffusers.py` to handle video generation using the Allegro inference pipeline.
- Implemented logging for better traceability during model initialization and video generation.
- Added error handling for robust pipeline setup and video export processes.
---
 .../allegro_settings.cpython-310.pyc          | Bin 0 -> 1693 bytes
 configs/allegro_settings.py                   |  34 +++++
 scripts/allegro_diffusers.py                  | 117 ++++++++++++++++++
 3 files changed, 151 insertions(+)
 create mode 100644 configs/__pycache__/allegro_settings.cpython-310.pyc
 create mode 100644 configs/allegro_settings.py
 create mode 100644 scripts/allegro_diffusers.py

diff --git a/configs/__pycache__/allegro_settings.cpython-310.pyc b/configs/__pycache__/allegro_settings.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..1c671650df5ef32ea7f02efdffc06dec74e6621e
GIT binary patch
literal 1693
zcmZux&2Aev5a#}9wNhlsiV`<SfP%*=P^5cc1$NUSm%t6;^s*O%UQ%netKAK`Ys8D(
zl53x!=smgiQF7$9C%-}tX@^|>*);{u47r@q@SAU#@9lLFj8CU&^0$xBKW4EzU=JO^
zR2T$9%t_JGi5v&LK+I*{N5nkteDETd`KM^yx`I<^?C9Q)RPj?@mLf}3Y&Lg?2a6+^
z>OBaKB8Q>K{S|^o9&;m~w|Kz3L&W^3ZS5A=K@{>1?6<)Vqb|P#b_eWk)U*3{!0ttR
z*4_iVAN8$0I5es5vmvB9{NP3dc9)GlqG+GpgM{zv!Eu`MMCR*c>n~jl;~2lXWHc*9
zjMeHII8Lcjc#%uI^1&ioaLKcn<2cV2BB>-TMV?K&HpZKCp;k_n;%XbQ%yA*}3&A+P
zps8Ti3Db;WPSqvm*@cjKw&Yp)<P30R*Q49mw`Wm^f~O+mo1CZia6OdgmrJgmP%&Gj
zuDi>e@swnA$+aI>jOu{#3la0H$50gQE6!O07L4zYbU&#eO}HpgF-^HZE}R)H={utM
z+sYt7Ql2H{`4s~PJ<6&j*?>b-$qRKgFkiey@qa!z@Z|`mdH`b0hcz3R7quF4)dx3f
zHGOVTuV2|)nzC(*cC8*emD|A<u$O|KrGQpd6uB%LuuK>o9>09~^NZKNkT)meK!?QQ
zM2PMZ0=QKv2fIhe+lr<JL0Gj3VR>whz+aF;@`ZS(N2~I30?ZMKE$L?@58#?sd>zkZ
zo|m(qs903LSF<K-HKPqD0o<m8rp2WWn#+z1P9ca8`Ht^}jx?af2Ez3)wgLM&O!W)|
zbKaxh(4Ww;pdHX2Xcx2(+5_DJ?Sl?Lw?MZ+2cW~h2GAs(`i~93uX$O?3{$L13B8Ra
zhpQ@JYQy1tfy?tP{cN)^epV%kIb2k-$e|Nf{b7)PDs8~0&~U4Ws#-3gFtu^mjkT>#
z{pbb{e%Nw(UUw$AA@aO#PoNPvo7ck$zE0qIeK5hd(AVzYLf&~DOfXP){?!?e+tR#>
zGBCnaLUxQW@2Kn<(FXx<Ahb}r*w$Kp2it~91~Nh(!b;r*@i2fo{%fb}fAih&Tm9h1
z{ibWDlkrgQgO3~<VZ5c`gUv^K0H7WaLbEJy%Avq-lcn}VR@zi`Ymw%(ypgfxk$hyr
zZTQ((CUI%rY%`Vl*ahzO+C69c>$|<SJ##Yl^srdHof`s6crTh5w*}YXVbfL540V-3
Jw}<Yz{{cFV$b$d?

literal 0
HcmV?d00001

diff --git a/configs/allegro_settings.py b/configs/allegro_settings.py
new file mode 100644
index 0000000..235376e
--- /dev/null
+++ b/configs/allegro_settings.py
@@ -0,0 +1,34 @@
+from pydantic_settings import BaseSettings
+
+class AllegroSettings(BaseSettings):
+    """
+    A Pydantic settings class for Allegro inference configuration.
+
+    This class uses Pydantic to provide validation and easy environment-based configuration
+    for Allegro inference pipeline settings.
+    """
+
+    model_name:str = "rhymes-ai/Allegro"
+    device: str = "cuda"
+    seed: int = 42
+    guidance_scale: float = 7.5
+    max_sequence_length: int = 512
+    num_inference_steps: int = 100
+    fps: int = 15
+
+    class Config:
+        """
+        Pydantic configuration class for environment variable support.
+        """
+        env_prefix = "ALLEGRO_"  # Prefix for environment variables
+        validate_assignment = True
+
+    def __repr__(self):
+        """
+        Return a string representation of the settings for debugging purposes.
+
+        :return: A string summarizing the settings.
+        """
+        return (f"AllegroSettings(model_name={self.model_name}, device={self.device}, seed={self.seed}, "
+                f"guidance_scale={self.guidance_scale}, max_sequence_length={self.max_sequence_length}, "
+                f"num_inference_steps={self.num_inference_steps}, fps={self.fps})")
diff --git a/scripts/allegro_diffusers.py b/scripts/allegro_diffusers.py
new file mode 100644
index 0000000..c1d5f29
--- /dev/null
+++ b/scripts/allegro_diffusers.py
@@ -0,0 +1,117 @@
+import torch
+from diffusers import AutoencoderKLAllegro, AllegroPipeline
+from diffusers.utils import export_to_video
+from loguru import logger
+from configs.allegro_settings import AllegroSettings
+
+class AllegroInference:
+    """
+    A class for managing the Allegro inference pipeline for generating videos based on textual prompts.
+
+    This class encapsulates the initialization, configuration, and video generation processes
+    for the Allegro model pipeline. It provides a streamlined way to handle prompts, model setup,
+    and output file management in a production-grade environment.
+    """
+
+    def __init__(self, settings: AllegroSettings):
+        """
+        Initialize the AllegroInference class with the given settings.
+
+        :param settings: An instance of AllegroSettings containing model, device, and generation parameters.
+        """
+        self.settings = settings
+        self.pipe = None
+
+        logger.info(f"Initializing {self.settings.model_name} inference pipeline")
+        self._setup_pipeline()
+
+    def _setup_pipeline(self):
+        """
+        Set up the Allegro model pipeline by loading the VAE and the pipeline with specified configurations.
+
+        This method loads the models, moves them to the specified device, and enables tiling for
+        efficient memory usage during inference.
+
+        :raises Exception: If there is an error during the model loading or configuration process.
+        """
+        try:
+            # Load VAE
+            logger.info("Loading VAE model...")
+            vae = AutoencoderKLAllegro.from_pretrained(
+                self.settings.model_name, 
+                subfolder="vae", 
+                torch_dtype=torch.float32
+            )
+
+            # Load Allegro pipeline
+            logger.info("Loading Allegro pipeline...")
+            self.pipe = AllegroPipeline.from_pretrained(
+                self.settings.model_name, 
+                vae=vae, 
+                torch_dtype=torch.bfloat16
+            )
+            
+            # Move pipeline to the specified device
+            self.pipe.to(self.settings.device)
+
+            # Enable tiling for efficient memory usage
+            self.pipe.vae.enable_tiling()
+
+            logger.info("Pipeline successfully initialized")
+        except Exception as e:
+            logger.error(f"Error initializing pipeline: {e}")
+            raise
+
+    def generate_video(self, prompt: str, positive_prompt: str, negative_prompt: str, output_path: str):
+        """
+        Generate a video based on the provided prompts and save it to the specified path.
+
+        :param prompt: The main textual description of the video scene.
+        :param positive_prompt: Additional positive prompts to enhance quality and style.
+        :param negative_prompt: Prompts to avoid undesirable features in the generated video.
+        :param output_path: File path to save the generated video.
+        :raises Exception: If there is an error during video generation or export.
+        """
+        try:
+            logger.info("Preparing prompts...")
+            prompt = positive_prompt.format(prompt.lower().strip())
+
+            logger.info("Generating video...")
+            generator = torch.Generator(device=self.settings.device).manual_seed(self.settings.seed)
+            video_frames = self.pipe(
+                prompt,
+                negative_prompt=negative_prompt,
+                guidance_scale=self.settings.guidance_scale,
+                max_sequence_length=self.settings.max_sequence_length,
+                num_inference_steps=self.settings.num_inference_steps,
+                generator=generator
+            ).frames[0]
+
+            logger.info(f"Exporting video to {output_path}...")
+            export_to_video(video_frames, output_path, fps=self.settings.fps)
+
+            logger.info("Video generation completed successfully")
+        except Exception as e:
+            logger.error(f"Error during video generation: {e}")
+            raise
+
+# Example usage (to be executed in a main script or testing environment)
+if __name__ == "__main__":
+    settings = AllegroSettings()
+    inference = AllegroInference(settings)
+
+    prompt = "A seaside harbor with bright sunlight and sparkling seawater, with many boats in the water. From an aerial view, the boats vary in size and color, some moving and some stationary. Fishing boats in the water suggest that this location might be a popular spot for docking fishing boats."
+    positive_prompt = """
+    (masterpiece), (best quality), (ultra-detailed), (unwatermarked), 
+    {} 
+    emotional, harmonious, vignette, 4k epic detailed, shot on kodak, 35mm photo, 
+    sharp focus, high budget, cinemascope, moody, epic, gorgeous
+    """
+
+    negative_prompt = """
+    nsfw, lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, 
+    low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry.
+    """
+
+    output_path = "output.mp4"
+    inference.generate_video(prompt, positive_prompt, negative_prompt, output_path)