複数replicateアカウント対応を修正

app.py

@@ -20,25 +20,24 @@ from PIL import Image as PIL_Image
 from jinja2 import Template
 
 ENV = os.getenv("ENV")
-MODEL = "gpt-3.5-turbo"
-# MODEL = "gpt-4"
+# MODEL = "gpt-3.5-turbo"
+MODEL = "gpt-4"
 
 load_dotenv()
 openai.api_key = os.getenv('OPENAI_API_KEY')
-# REPLICATE_API_TOKEN = os.getenv("REPLICATE_API_TOKEN")
-# REPLICATE_API_TOKEN_LISTをロードし、カンマで分割してリストに変換
 REPLICATE_API_TOKEN_LIST = os.getenv("REPLICATE_API_TOKEN_LIST").split(',')
-REPLICATE_API_TOKEN_INDEX = 0 # トークンのインデックスを初期化
 NUMBER_OF_SCENES = os.getenv("NUMBER_OF_SCENES")
 
 if ENV == "PRODUCTION":
     import replicate
+    from replicate.client import Client
 else:
-    from stub import replicate
+    # from stub import replicate
+    pass
 
 class Video:
-    def __init__(self, scene, index, token_controller):
-        self.token_controller = token_controller
+    def __init__(self, scene, index, client: Client):
+        self.client = client
         self.scene = scene
         self.prompt = "masterpiece, awards, best quality, dramatic-lighting, "
         self.prompt = self.prompt + scene.get("visual_prompt_in_en")
@@ -49,40 +48,77 @@ class Video:
         self.video_id = uuid.uuid4()
         self.file_path = f"assets/thread_{index}_request_{self.video_id}_video.mp4"
 
-    MAX_RETRIES = 2
+    REPLICATE_MODEL_PATH = "lucataco/animate-diff"
+    REPLICATE_MODEL_VERSION = "1531004ee4c98894ab11f8a4ce6206099e732c1da15121987a8eef54828f0663"
+
     def run_replicate(self, retries=0):
         try:
-            self.token = self.token_controller.get_next_token()
+            # self.client.api_token = self.client.api_token_controller.get_next_token()
             start_time = time.time()
 
-            os.environ["REPLICATE_API_TOKEN"] = self.token
+            # os.environ["REPLICATE_API_TOKEN"] = self.client.api_token
             #tokenの最初の10文字だけ出力
-            print(f"Thread {self.index} token: {self.token[:10]}")
+            print(f"Thread {self.index} token: {self.client.api_token[:10]}")
             
-            self.output_url = replicate.run(
-                "lucataco/animate-diff:1531004ee4c98894ab11f8a4ce6206099e732c1da15121987a8eef54828f0663",
+            model = self.client.models.get(self.REPLICATE_MODEL_PATH)
+            version = model.versions.get(self.REPLICATE_MODEL_VERSION)
+            self.prediction = self.client.predictions.create(
+                version=version,
                 input={
                     "motion_module": "mm_sd_v14", 
                     "prompt": self.prompt,
                     "n_prompt": self.nagative_prompt,
                     "seed": 0,
-                    }
+                    },
             )
+            
+            self.prediction_id = self.prediction.id
+
+            # print(f"Thread {self.index} token: {self.client.api_token[:10]} prediction: {self.prediction}")
+            print(f"Thread {self.index} token: {self.client.api_token[:10]} prediction.status: {self.prediction.status}")
+
+            self.prediction.reload()
+            print(f"Thread {self.index} token: {self.client.api_token[:10]} prediction.status: {self.prediction.status}")
 
+            self.prediction.wait()
+
+            print(f"Thread {self.index} token: {self.client.api_token[:10]} prediction.status: {self.prediction.status}")
+            if self.prediction.status == "succeeded":
+                self.output_url = self.prediction.output
+                print(f"Thread {self.index} token: {self.client.api_token[:10]} prediction.output: {self.prediction.output}")
+            else:
+                self.output_url = None
+            
             end_time = time.time()
             duration = end_time - start_time
 
             self.download_and_save(url=self.output_url, file_path=self.file_path)
             self.print_thread_info(start_time, end_time, duration)
         except replicate.exceptions.ReplicateError as e:
-            if str(e) == "The requested resource could not be found." and retries < self.MAX_RETRIES:
-                print("リソースが見つからないエラーが発生しました。2秒後に再試行します。")
-                time.sleep(2)
-                self.run_replicate(retries + 1)  # 再帰的に関数を呼び出して再試行
-            elif retries >= self.MAX_RETRIES:
-                print("最大再試行回数に達しました。スレッドを終了します。")
-                # 最大再試行回数に達した場合の追加処理
+            if self.prediction and str(e) == "The requested resource could not be found.":
+                predictions = self.client.predictions.list()
+                self.prediction = next((p for p in predictions if p.id == self.prediction_id), None)
+
+                if self.prediction:
+                    print(f"Found prediction with ID {self.prediction_id}: {self.prediction}")
+                else:
+                    print(f"No prediction found with ID {self.prediction_id}")
+                self.prediction.wait()
+
+                print(f"Thread {self.index} token: {self.client.api_token[:10]} prediction.status: {self.prediction.status}")
+                if self.prediction.status == "succeeded":
+                    self.output_url = self.prediction.output
+                    print(f"Thread {self.index} token: {self.client.api_token[:10]} prediction.output: {self.prediction.output}")
+                else:
+                    self.output_url = None
+                
+                end_time = time.time()
+                duration = end_time - start_time
+
+                self.download_and_save(url=self.output_url, file_path=self.file_path)
+                self.print_thread_info(start_time, end_time, duration)
             else:
+                print(f"Error in thread {self.index}: {e}")
                 print("予期しないエラーが発生しました。スレッドを終了します。")
                 # 予期しないエラーが発生した場合の追加処理
         except Exception as e:
@@ -105,24 +141,37 @@ class ThreadController:
         scenes = args.get("scenes")
         self.videos = []
         self.threads = []
-        self.token_index = 0
         self.lock = threading.Lock()
-        for index, scene in enumerate(scenes):
-            for _ in REPLICATE_API_TOKEN_LIST:
-                # token = REPLICATE_API_TOKEN_LIST[self.token_index]
-                video = Video(scene, index, self)
+        self.replicate_client_list = {}
+        for token in REPLICATE_API_TOKEN_LIST:
+
+            client = Client()
+            client.api_token = token
+            self.replicate_client_list[token] = client
+            
+            for index, scene in enumerate(scenes):
+                # token = REPLICATE_API_TOKEN_LIST[self.client.api_token_index]
+                video = Video(scene, index, client)
                 self.videos.append(video)
-                self.token_index = (self.token_index + 1) % len(REPLICATE_API_TOKEN_LIST)
+            
+            # self.client.api_token_index = (self.client.api_token_index + 1) % len(REPLICATE_API_TOKEN_LIST)
+            
 
     def run_threads(self):
         os.makedirs("assets", exist_ok=True)
 
+        token = None
         for video in self.videos:
+            if token is not None and video.client.api_token != token:
+                # tokenが異なる場合、1秒待ってから次を実行
+                print(f"Thread {video.index} token changed. Waiting 3 seconds.")
+                time.sleep(5)
+            
             thread = threading.Thread(target=video.run_replicate)
             self.threads.append(thread)
             thread.start()
-            # 1秒待ってから実行
-            # time.sleep(1)
+            token = video.client.api_token
+            # time.sleep(5)
 
         for thread in self.threads:
             thread.join()
@@ -150,12 +199,6 @@ class ThreadController:
         for video in self.videos:
             print(f"Thread {video.index} prompt: {video.prompt}")
 
-    def get_next_token(self):
-        with self.lock:
-            token = REPLICATE_API_TOKEN_LIST[self.token_index]
-            self.token_index = (self.token_index + 1) % len(REPLICATE_API_TOKEN_LIST)
-            return token
-
 def main(args):
     thread_controller = ThreadController(args)
     thread_controller.run_threads()
@@ -213,34 +256,6 @@ class NajiminoAI:
 
     def generate_markdown(self, args, generation_time):
 
-        # # lang=args.get("lang")
-        # title=args.get("title")
-        # description=args.get("description")
-        # visual_prompt_in_en=args.get("visual_prompt_in_en")
-        # scenes = args.get("scenes")
-        
-        # prompt_for_visual_expression = \
-        #     visual_prompt_in_en
-            
-        # print("prompt_for_visual_expression: "+prompt_for_visual_expression)
-        
-        # prompts = []
-        # if scenes:
-        #     for scene_data in scenes:
-        #         prompt = scene_data.get("visual_prompt_in_en")
-        #         prompt = prompt + ", " + scene_data.get("cinematic_angles")
-        #         prompt = prompt + ", " + scene_data.get("visual_prompt_in_en")
-        #         prompts.append(prompt)
-                # print("scenes: " + json.dumps(scenes, indent=2))
-        # if scenes:
-        #     for scene_data in scenes:
-        #         scene = scene_data.get("scene")
-        #         cinematic_angles = scene_data.get("cinematic_angles")
-        #         visual_prompt_in_en = scene_data.get("visual_prompt_in_en")
-        #         print("scene: ", scene)
-        #         print("cinematic_angles: ", cinematic_angles)
-        #         print("visual_prompt_in_en: ", visual_prompt_in_en)
-
         template_string = get_filetext(filename = "template.md")
         
         template = Template(template_string)
@@ -282,29 +297,15 @@ class NajiminoAI:
             
         function_name = message["function_call"]["name"]
         
-        args = json.loads(message["function_call"]["arguments"])
+        try:
+            args = json.loads(message["function_call"]["arguments"])
+        except json.JSONDecodeError as e:
+            print(f"JSON decode error at position {e.pos}: {e.msg}")
+            print("message: " + json.dumps(message, indent=2))
+            raise e
         
         print("args: " + json.dumps(args, indent=2))
         
-        # # lang=args.get("lang")
-        # title=args.get("title")
-        # description=args.get("description")
-        # visual_prompt_in_en=args.get("visual_prompt_in_en")
-        # scenes = args.get("scenes")
-        
-        # prompt_for_visual_expression = \
-        #     visual_prompt_in_en
-            
-        # print("prompt_for_visual_expression: "+prompt_for_visual_expression)
-        
-        # prompts = []
-        # if scenes:
-        #     for scene_data in scenes:
-        #         prompt = scene_data.get("visual_prompt_in_en")
-        #         prompt = prompt + ", " + scene_data.get("cinematic_angles")
-        #         prompt = prompt + ", " + scene_data.get("visual_prompt_in_en")
-        #         prompts.append(prompt)
-        
         video_path = main(args)
             
         main_end_time = time.time()
@@ -325,7 +326,6 @@ class NajiminoAI:
         )
         return [video_path, html]
 
-
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="Generate videos from text prompts")
 
@@ -338,10 +338,6 @@ if __name__ == "__main__":
         # main(prompts)
         NajiminoAI.generate("伝統工芸と最新技術の融合")
     else:
-        # def create_video(prompt):
-        #     prompts = prompt.strip().split('\n')
-        #     output_path = main(prompts)
-        #     return output_path
 
         iface = gr.Interface(
             fn=NajiminoAI.generate,
@@ -357,7 +353,7 @@ if __name__ == "__main__":
             examples=[
                 ["侍たちは野を超え山を超え、敵軍大将を討ち取り、天下の大将軍となった！"],
                 ["子どもたちが笑ったり怒ったり泣いたり楽しんだりする"],
-                ["彼女のダンスは炎のように激しく、風のように自由に、水のように柔軟に、木のように生き生きと、虹のように美しく舞う"],
+                ["日は昇り、大地を照らし、日は沈む。闇夜を照らし、陽はまた昇る。 "],
             ],
         )
         iface.launch()

schema.json

@@ -20,7 +20,7 @@
 					},
 					"story": {
 						"type": "string",
-						"description": "映画のあらすじを、起承転結を交えて時系列に詳しく説明する"
+						"description": "映画のあらすじを、起承転結を交えて時系列に文学的に詳しく説明する"
 					},
 					"visual_style": {
 						"type": "string",

template.md

@@ -12,4 +12,4 @@
 
 | Scene | visual_prompt_in_en | negative_visual_prompt_in_en | cinematic_angles |
 |----:|----|----|----|{% for item in args.scenes %}
-|{{ item.scene }}|{{ item.visual_prompt_in_en }}|{{ item.negative_visual_prompt_in_en }}|{{ item.camera_work}}|{% endfor %}
+|{{ item.scene }}|{{ item.visual_prompt_in_en }}|{{ item.negative_visual_prompt_in_en }}|{{ item.cinematic_angles}}|{% endfor %}