Merge pull request #107 from banodoco/green-head

Green head

.streamlit/credentials.toml

@@ -0,0 +1,2 @@
+[general]
+email=""

app.py

@@ -50,8 +50,10 @@ def start_runner():
             app_logger.info("Starting runner")
             python_executable = sys.executable
             _ = subprocess.Popen([python_executable, "banodoco_runner.py"])
-            while not is_process_active(RUNNER_PROCESS_NAME, RUNNER_PROCESS_PORT):
+            max_retries = 6
+            while not is_process_active(RUNNER_PROCESS_NAME, RUNNER_PROCESS_PORT) and max_retries:
                 time.sleep(0.1)
+                max_retries -= 1
         else:
             # app_logger.debug("Runner already running")
             pass

backend/db_repo.py

@@ -194,6 +194,7 @@ def get_all_file_list(self, **kwargs):
 
             kwargs['project_id'] = project.id
 
+        # hackish sol: you can pass custom params as long as 'page' is not in the kwargs
         if 'page' in kwargs and kwargs['page']:
             page = kwargs['page']
             del kwargs['page']

backend/models.py

@@ -321,6 +321,10 @@ def save(self, *args, **kwargs):
             self.aux_frame_index = new_index
             timing_list.update(aux_frame_index=F('aux_frame_index') - 1)
 
+        # --------------- adding alternative images ----------
+        if not (self.alternative_images and len(self.alternative_images)) and self.primary_image:
+            self.alternative_images = json.dumps([str(self.primary_image.uuid)])
+
         super().save(*args, **kwargs)
 
 

banodoco_runner.py

@@ -21,7 +21,7 @@
 from ui_components.methods.file_methods import get_file_bytes_and_extension, load_from_env, save_or_host_file_bytes, save_to_env
 from utils.common_utils import acquire_lock, release_lock
 from utils.data_repo.data_repo import DataRepo
-from utils.ml_processor.constants import replicate_status_map
+from utils.ml_processor.constants import ComfyWorkflow, replicate_status_map
 
 from utils.constants import RUNNER_PROCESS_NAME, RUNNER_PROCESS_PORT, AUTH_TOKEN, REFRESH_AUTH_TOKEN
 from utils.ml_processor.gpu.utils import is_comfy_runner_present, predict_gpu_output, setup_comfy_runner
@@ -75,7 +75,7 @@ def main():
         server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
         server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
         server_socket.bind(("localhost", RUNNER_PROCESS_PORT))
-        server_socket.listen(1)
+        server_socket.listen(100)     # hacky fix
 
     print('runner running')
     while True:
@@ -176,13 +176,19 @@ def find_process_by_port(port):
 
         return pid
 
-def stop_server(self, port):
+def stop_server(port):
         pid = find_process_by_port(port)
         if pid:
             app_logger.log(LoggingType.DEBUG, "comfy server stopped")
             process = psutil.Process(pid)
             process.terminate()
             process.wait()
+
+def format_model_output(output, model_display_name):
+    if model_display_name and model_display_name == ComfyWorkflow.MOTION_LORA.value:
+        return output
+    else:
+        return [output[-1]]
 
 def check_and_update_db():
     # print("updating logs")
@@ -308,11 +314,15 @@ def check_and_update_db():
                     data['output_node_ids'], data.get("extra_model_list", []), data.get("ignore_model_list", []))
                 end_time = time.time()
 
-                output = output[-1]     # TODO: different models can have different logic
-                destination_path = "./videos/temp/" + str(uuid.uuid4()) + "." + output.split(".")[-1]
-                shutil.copy2("./output/" + output, destination_path)
+                res_output = format_model_output(output, log.model_name)
+                destination_path_list = []
+                for output in res_output:
+                    destination_path = "./videos/temp/" + str(uuid.uuid4()) + "." + output.split(".")[-1]
+                    shutil.copy2("./output/" + output, destination_path)
+                    destination_path_list.append(destination_path)
+
                 output_details = json.loads(log.output_details)
-                output_details['output'] = destination_path
+                output_details['output'] = destination_path_list[0] if len(destination_path_list) == 1 else destination_path_list
                 update_data = {
                     "status" : InferenceStatus.COMPLETED.value,
                     "output_details" : json.dumps(output_details),
@@ -321,7 +331,7 @@ def check_and_update_db():
 
                 InferenceLog.objects.filter(id=log.id).update(**update_data)
                 origin_data = json.loads(log.input_params).get(InferenceParamType.ORIGIN_DATA.value, {})
-                origin_data['output'] = destination_path
+                origin_data['output'] = destination_path_list[0] if len(destination_path_list) == 1 else destination_path_list
                 origin_data['log_uuid'] = log.uuid
                 print("processing inference output")
 

banodoco_settings.py

@@ -98,7 +98,9 @@ def create_new_user_data(user: InternalUserObject):
 def create_new_project(user: InternalUserObject, project_name: str, width=512, height=512):
     data_repo = DataRepo()
 
-    # creating a new project for this user
+    existing_projects_list = data_repo.get_all_project_list(user.uuid)
+    add_initial_frames = False if (existing_projects_list and len(existing_projects_list)) else True
+
     project_data = {
         "user_id": user.uuid,
         "name": project_name,
@@ -115,43 +117,46 @@ def create_new_project(user: InternalUserObject, project_name: str, width=512, h
     }
 
     shot = data_repo.create_shot(**shot_data)
-
-    # create timings for init_images 
-    init_images_path = os.path.join("sample_assets", "sample_images", "init_frames")
-    init_image_list = list_files_in_folder(init_images_path)
     st.session_state["project_uuid"] = project.uuid
-
-    for idx, img_path in enumerate(init_image_list):
-        img_path = os.path.join(init_images_path, img_path)
-        img = Image.open(img_path)
-        img = img.resize((width, height))
-
-        unique_file_name = f"{str(uuid.uuid4())}.png"
-        file_location = f"videos/{project.uuid}/resources/prompt_images/{unique_file_name}"
-        hosted_url = save_or_host_file(img, file_location, mime_type='image/png', dim=(width, height))
-        file_data = {
-            "name": str(uuid.uuid4()),
-            "type": InternalFileType.IMAGE.value,
-            "project_id": project.uuid,
-            "dim": (width, height),
-        }
-
-        if hosted_url:
-            file_data.update({'hosted_url': hosted_url})
-        else:
-            file_data.update({'local_path': file_location})
 
-        source_image = data_repo.create_file(**file_data)
-
-        timing_data = {
-            "frame_time": 0.0,
-            "aux_frame_index": idx,
-            "source_image_id": source_image.uuid,
-            "shot_id": shot.uuid,
-        }
-        timing: InternalFrameTimingObject = data_repo.create_timing(**timing_data)
-
-        add_image_variant(source_image.uuid, timing.uuid)
+    # Add initial frames only if there are no existing projects (i.e., it's the user's first project)
+    if add_initial_frames:
+        init_images_path = os.path.join("sample_assets", "sample_images", "init_frames")
+        init_image_list = list_files_in_folder(init_images_path)
+        image_extensions = {'.png', '.jpg', '.jpeg', '.gif'}
+        init_image_list = [img for img in init_image_list if os.path.splitext(img)[1].lower() in image_extensions]
+
+        for idx, img_path in enumerate(init_image_list):
+            img_path = os.path.join(init_images_path, img_path)
+            img = Image.open(img_path)
+            img = img.resize((width, height))
+
+            unique_file_name = f"{str(uuid.uuid4())}.png"
+            file_location = f"videos/{project.uuid}/resources/prompt_images/{unique_file_name}"
+            hosted_url = save_or_host_file(img, file_location, mime_type='image/png', dim=(width, height))
+            file_data = {
+                "name": str(uuid.uuid4()),
+                "type": InternalFileType.IMAGE.value,
+                "project_id": project.uuid,
+                "dim": (width, height),
+            }
+
+            if hosted_url:
+                file_data.update({'hosted_url': hosted_url})
+            else:
+                file_data.update({'local_path': file_location})
+
+            source_image = data_repo.create_file(**file_data)
+
+            timing_data = {
+                "frame_time": 0.0,
+                "aux_frame_index": idx,
+                "source_image_id": source_image.uuid,
+                "shot_id": shot.uuid,
+            }
+            timing: InternalFrameTimingObject = data_repo.create_timing(**timing_data)
+
+            add_image_variant(source_image.uuid, timing.uuid)
 
     # create default ai models
     model_list = create_predefined_models(user)

readme.md

@@ -97,15 +97,15 @@ Then go to this URL, and it should be running!
 
 This commands sets up the app. Run this only the first time, after that you can simply start the app using the next command.
 ```bash
-curl -sSL https://raw.githubusercontent.com/banodoco/Dough/green-head/linux_setup.sh | bash
+curl -sSL https://raw.githubusercontent.com/banodoco/Dough/green-head/scripts/linux_setup.sh | bash
 ```
 
 ### Run the app
 
 you can run the app using 
 
 ```bash
-source ./dough-env/bin/activate && sh entrypoint.sh
+source ./dough-env/bin/activate && sh ./scripts/entrypoint.sh
 ```
 
 ## Instructions for Windows:
@@ -117,14 +117,14 @@ source ./dough-env/bin/activate && sh entrypoint.sh
 Run the setup script
 
 ```bash
-iwr -useb "https://raw.githubusercontent.com/banodoco/Dough/green-head/windows_setup.bat" -OutFile "script.bat"
+iwr -useb "https://raw.githubusercontent.com/banodoco/Dough/green-head/scripts/windows_setup.bat" -OutFile "script.bat"
 Start-Process "cmd.exe" -ArgumentList "/c script.bat"
 ```
 
 ### Run the app
 
 ```bash
-. .\dough-env\Scripts\activate ; .\entrypoint.bat
+. .\dough-env\Scripts\activate ; .\scripts\entrypoint.bat
 ```
 
 If you're having any issues, please share them in our [Discord](https://discord.com/invite/8Wx9dFu5tP).