Merge branch 'main' of github.com:Dungeon-MASSters/MASSter

model/generator/main.py

@@ -6,13 +6,14 @@
 from libs import logger
 from libs.models import ImageGenerationModel
 import urllib.request
+from libs.config import settings
 
 class Pipeline:
     def __init__(self) -> None:
-        self.pb: PocketBase = PocketBase("https://pb.apps.npod.space/")
+        self.pb: PocketBase = PocketBase(settings.PB_LINK)
         self.admin_data = self.pb.admins.auth_with_password(
-            "[email protected]", 
-            "6c6297287af76472")
+            settings.PB_LOGIN, 
+            settings.PB_PWD)
 
         self.model = ImageGenerationModel()
         self.model.load_kandinsky()
@@ -21,7 +22,7 @@ def run(self):
         while True:
            # logger.info('listening to the collection')
             try:
-                records = self.pb.collection('text_generation_mvp').get_list(
+                records = self.pb.collection(settings.COL_NAME).get_list(
                     page=1,
                     per_page=1,
                     query_params={
@@ -47,7 +48,7 @@ def run(self):
 
                 images = self.process_generation(records.items[0], width=width, height=height)
                 if len(images) == 0:
-                    self.pb.collection('text_generation_mvp').update(
+                    self.pb.collection(settings.COL_NAME).update(
                         id=records.items[0].id,
                         body_params={
                             "status": "error"
@@ -56,7 +57,7 @@ def run(self):
                     logger.error('Не было создано ни одного изображения')
                     continue
 
-                self.pb.collection('text_generation_mvp').update(
+                self.pb.collection(settings.COL_NAME).update(
                     id=records.items[0].id,
                     body_params={
                         "status": "generated",
@@ -74,7 +75,7 @@ def run(self):
                     image.save(path)
                     uploads.append((path, open(path, mode='rb')))
 
-                self.pb.collection('text_generation_mvp').update(
+                self.pb.collection(settings.COL_NAME).update(
                     id=records.items[0].id,
                     body_params={
                         "status": "generated",
@@ -92,7 +93,7 @@ def process_generation(self, data, width, height) -> list[Image.Image]:
 
         if data.input_image != '' and data.type == 'avatar':
             try:
-                image_url = self.pb.collection('text_generation_mvp').get_file_url(data, filename=data.input_image) 
+                image_url = self.pb.collection(settings.COL_NAME).get_file_url(data, filename=data.input_image) 
                 urllib.request.urlretrieve(image_url, data.input_image)
 
                 self.model.remove_models()

model/preprocessor/libs/config.py

@@ -0,0 +1,12 @@
+from pydantic_settings import BaseSettings, SettingsConfigDict
+
+class Settings(BaseSettings):
+    model_config = SettingsConfigDict(env_file='.env', env_file_encoding='utf-8', extra='ignore')
+
+class ModuleSettings(Settings):
+    PB_LINK: str = "https://pb.apps.npod.space/"
+    PB_LOGIN: str = "[email protected]"
+    PB_PWD: str = "6c6297287af76472"
+    COL_NAME: str = "text_generation_mvp"
+
+settings = ModuleSettings()

model/preprocessor/libs/models.py

@@ -100,20 +100,30 @@ def video_parser(self, video_file):
             video_clip.save_frame(frame_filename, current_duration)
         return folder_dir
 
-    def get_desc_from_video(self, video_file: str) -> str:
-        if not video_file.endswith('.mp4'):
-            return ''
+    def get_desc_from_video(self, video_files) -> str:
+
+        result_summary = []
+        for video_file in video_files:
+            if not video_file.endswith('.mp4'):
+                return ''
 
-        folder_dir = self.video_parser(video_file)
-        desc = []
-        for images in os.listdir(folder_dir):
-            desc.append(self.generate_desc(folder_dir + '/' + images))
+            folder_dir = self.video_parser(video_file)
+            desc = []
+            for images in os.listdir(folder_dir):
+                desc.append(self.generate_desc(folder_dir + '/' + images))
 
-        if os.path.exists(folder_dir):
-            shutil.rmtree(folder_dir)
+            if os.path.exists(folder_dir):
+                shutil.rmtree(folder_dir)
 
-        summary_desc = self.get_summary(desc)
-        return summary_desc
+            summary_desc = self.get_summary(desc)
+            result_summary.append(summary_desc)
+
+        if len(result_summary) > 1:
+            return self.get_summary(result_summary)
+        elif len(result_summary) == 1:
+            return result_summary[0]
+        else:
+            return ''
 
     def get_summary(self, desc) -> str:
 

model/preprocessor/main.py

@@ -9,21 +9,23 @@
 from libs import logger
 from libs.models import Translator, VideoParser
 import urllib.request
+from libs.config import settings
+import re
 
 class Pipeline:
     def __init__(self) -> None:
-        self.pb: PocketBase = PocketBase("https://pb.apps.npod.space/")
+        self.pb: PocketBase = PocketBase(settings.PB_LINK)
         self.admin_data = self.pb.admins.auth_with_password(
-            "[email protected]", 
-            "6c6297287af76472")
+            settings.PB_LOGIN, 
+            settings.PB_PWD)
 
         self.translator = Translator()
         self.vparser = VideoParser()
 
     def run(self):
         while True:
             try:
-                records = self.pb.collection('text_generation_mvp').get_list(
+                records = self.pb.collection(settings.COL_NAME).get_list(
                     page=1,
                     per_page=1,
                     query_params={
@@ -53,12 +55,21 @@ def run(self):
                     style = records.items[0].style
 
                 try:
-                    video_url = self.pb.collection('text_generation_mvp').get_file_url(records.items[0], filename=records.items[0].video)
+                    video_url = self.pb.collection(settings.COL_NAME).get_file_url(records.items[0], filename=records.items[0].video)
 
-                    video_name = 'video.mp4' 
-                    urllib.request.urlretrieve(video_url, video_name)
-
-                    summary = self.vparser.get_desc_from_video(video_name)[0]['summary_text']
+                    s = re.findall('\[(.*)\]', url)
+                    ss = re.sub('[\' ]', '', s[0])
+                    video_names = re.split(',', ss)
+
+                    url = re.sub('\[.*\]', '', url)
+
+                    #video_name = 'video.mp4' 
+
+                    for video_name in video_names:
+                        urllib.request.urlretrieve(video_url + video_name, video_name)
+
+
+                    summary = self.vparser.get_desc_from_video(video_names)[0]['summary_text']
                     print(summary)
 
                     if prompt != '':
@@ -68,7 +79,7 @@ def run(self):
                 except:
                     logger.info('Нет приложенного видео')
 
-                self.pb.collection('text_generation_mvp').update(
+                self.pb.collection(settings.COL_NAME).update(
                     id=records.items[0].id,
                     body_params={
                         "status": "preprocessed",