Update app.py

app.py

@@ -15,8 +15,8 @@ from TTS.tts.models.xtts import Xtts
 from vinorm import TTSnorm
 from langchain.llms import HuggingFacePipeline
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
-from components import caption_chain, tag_chain
-from components import pexels, utils
+from diffusers import StableDiffusionPipeline
+from PIL import Image
 import cv2
 from moviepy.editor import AudioFileClip, ImageSequenceClip
 import gc
@@ -73,6 +73,12 @@ llm_chain = caption_chain.chain(llm=local_llm)
 sum_llm_chain = tag_chain.chain(llm=local_llm)
 pexels_api_key = os.getenv('pexels_api_key')
 
+# Initialize Stable Diffusion Pipeline with TDN-M/East-asian-beauty
+image_gen_model_id = "TDN-M/East-asian-beauty"
+device = "cuda" if torch.cuda.is_available() else "cpu"
+image_generator = StableDiffusionPipeline.from_pretrained(image_gen_model_id, torch_dtype=torch.float16)
+image_generator = image_generator.to(device)
+
 def normalize_vietnamese_text(text):
     text = (
         TTSnorm(text, unknown=False, lower=False, rule=True)
@@ -101,7 +107,7 @@ def calculate_keep_len(text, lang):
         return 13000 * word_count + 2000 * num_punct
     return -1
 
-def create_video_from_audio(audio_path, images, output_path):
+def create_video_from_audio_and_images(audio_path, images, output_path):
     audio_clip = AudioFileClip(audio_path)
     duration = audio_clip.duration
     
@@ -128,6 +134,19 @@ def truncate_prompt(prompt, tokenizer, max_length=512):
         prompt = tokenizer.convert_tokens_to_string(tokens)
     return prompt
 
+def generate_images_from_sentences(sentences, image_generator, folder_path):
+    try:
+        for i, sentence in enumerate(sentences):
+            print(f"Generating image for sentence {i + 1}: {sentence}")
+            image = image_generator(sentence, guidance_scale=7.5).images[0]
+            image_path = os.path.join(folder_path, f"image_{i + 1}.png")
+            image.save(image_path)
+            print(f"Saved image at {image_path}")
+    except Exception as e:
+        print("Error! Failed generating images")
+        print(e)
+        return []
+
 @spaces.GPU
 def predict(
     prompt,
@@ -207,18 +226,23 @@ def predict(
         out["wav"] = out["wav"][:keep_len]
         torchaudio.save("output.wav", torch.tensor(out["wav"]).unsqueeze(0), 24000)
         
-        # Tạo video từ file audio
-        print("I: Generating video from audio...")
+        # Tạo video từ file audio và các cảnh
+        print("I: Generating images from sentences...")
         # Sử dụng UUID để tạo tên thư mục ngắn gọn
         folder_name = f"video_{uuid.uuid4().hex}"
         os.makedirs(folder_name, exist_ok=True)
         folder_path = os.path.join(folder_name, "images")
         os.makedirs(folder_path, exist_ok=True)
         
-        # Tạo video từ file audio và các hình ảnh
-        folder_name, sentences = pexels.generate_videos(prompt, pexels_api_key, "landscape", 1080, 1920, llm_chain, sum_llm_chain)
-        utils.combine_videos(folder_name)
+        # Tách các câu từ văn bản
+        sentences = [x.strip() for x in re.split(r'[.!?]', prompt) if len(x.strip()) > 6]
+        
+        # Tạo ảnh minh họa cho từng câu
+        images = generate_images_from_sentences(sentences, image_generator, folder_path)
+        
+        # Tạo video từ file audio và các ảnh
         video_path = os.path.join(folder_name, "Final_Ad_Video.mp4")
+        create_video_from_audio_and_images("output.wav", images, video_path)
         
         print(f"I: Video generated at {video_path}")
         metrics_text += f"Video generated at {video_path}\n"