Update app.py

app.py

@@ -10,6 +10,9 @@ import os
 from transformers import pipeline
 import cv2
 import io
+import json
+import re
+import time
 
 detector = dlib.get_frontal_face_detector()
 try:
@@ -24,36 +27,58 @@ except RuntimeError:
         f.write(landmarks_data)
     predictor = dlib.shape_predictor("shape_predictor_68_face_landmarks.dat")
 
-API_URL = "https://api-inference.huggingface.co/models/black-forest-labs/flux-1-schnell"
+IMAGE_GEN_API = "https://api-inference.huggingface.co/models/black-forest-labs/flux-1-schnell"
 HF_TOKEN = os.getenv("HF_TOKEN")
 
-LLM_API_URL = "https://api-inference.huggingface.co/models/lmsys/fastchat-t5-3b-v1.0"
+LLM_API = "https://api-inference.huggingface.co/models/lmsys/fastchat-t5-3b-v1.0"
 
 def query_hf_image_generation(prompt):
     headers = {"Authorization": f"Bearer {HF_TOKEN}"}
     payload = {"inputs": prompt}
-    response = requests.post(API_URL, headers=headers, json=payload)
-    if response.status_code == 200:
-        image_bytes = response.content
-        image = Image.open(io.BytesIO(image_bytes))
-        return image
-    else:
-        raise Exception(f"Image generation failed: {response.content}")
-
-def query_llm(prompt, image_description):
+    for _ in range(3):
+        try:
+            response = requests.post(IMAGE_GEN_API, headers=headers, json=payload)
+            response.raise_for_status()
+            image_bytes = response.content
+            image = Image.open(io.BytesIO(image_bytes))
+            return image
+        except requests.exceptions.RequestException as e:
+            print(f"Image generation attempt failed: {e}")
+            time.sleep(2)
+    raise Exception("Image generation failed after multiple attempts")
+
+def query_llm(prompt, system_prompt):
     headers = {"Authorization": f"Bearer {HF_TOKEN}"}
-    system_prompt = "You are an expert in image to video creation, and give only the motion type, intensity, text overlay, text color, text start and end times for the image described below based on user's prompt. Give the response in a JSON format."
-    prompt_template = f"<|system|>\n{system_prompt}</s>\n<|user|>\nImage Description: {image_description}\nUser Prompt: {prompt}</s>\n<|assistant|>\n"
+    prompt_template = f"<|system|>\n{system_prompt}</s>\n<|user|>\n{prompt}</s>\n<|assistant|>\n"
     payload = {"inputs": prompt_template, "max_new_tokens": 200}
-    response = requests.post(LLM_API_URL, headers=headers, json=payload)
-    if response.status_code == 200:
-        return response.json()[0]['generated_text']
-    else:
-        raise Exception(f"LLM query failed: {response.content}")
+    for _ in range(3):
+        try:
+            response = requests.post(LLM_API, headers=headers, json=payload)
+            response.raise_for_status()
+            return response.json()[0]['generated_text']
+        except requests.exceptions.RequestException as e:
+            print(f"LLM query attempt failed: {e}")
+            time.sleep(2)
+    raise Exception("LLM query failed after multiple attempts")
+
+def segment_script(script):
+    system_prompt = "You are a helpful assistant. Given a script, divide it into segments suitable for generating images, ensuring each segment is less than 500 characters."
+    llm_response = query_llm(script, system_prompt)
+    segments = llm_response.split('\n')
+    segments = [seg.strip() for seg in segments if seg.strip()]
+    return segments
+
+def generate_image_prompts(script_segments):
+    image_prompts = []
+    for segment in script_segments:
+        system_prompt = "You are a helpful assistant. Create a concise image prompt based on the following script segment:"
+        prompt = f"Script Segment: {segment}"
+        image_prompt = query_llm(prompt, system_prompt)
+        image_prompts.append(image_prompt)
+    return image_prompts
 
 def extract_motion_params(llm_output):
     try:
-        import json
         start_index = llm_output.find('{')
         end_index = llm_output.rfind('}') + 1
         json_string = llm_output[start_index:end_index]
@@ -78,7 +103,7 @@ def detect_face_landmarks(image):
         return shape
     else:
         return None
-    
+
 def apply_color_grading(frame, color_preset, intensity):
     if color_preset == "sepia":
         sepia_matrix = np.array([[0.393, 0.769, 0.189],
@@ -123,7 +148,7 @@ def apply_color_grading(frame, color_preset, intensity):
         return (np.clip(frame_float, 0, 1) * 255).astype(np.uint8)
 
     return frame
-    
+
 def apply_vignette(frame, intensity):
     width, height = frame.shape[1], frame.shape[0]
     x = np.linspace(-1, 1, width)
@@ -201,7 +226,7 @@ def apply_advanced_motion(image, motion_type, intensity, duration, fps, text_ove
                     draw.line([tuple(landmarks[42]), tuple(landmarks[45])], fill=text_color, width=2)
                 else:
                     frame = image.copy()
-            
+
             elif motion_type == "smile":
                 mouth_left = landmarks[48]
                 mouth_right = landmarks[54]
@@ -314,14 +339,53 @@ def apply_advanced_motion(image, motion_type, intensity, duration, fps, text_ove
 
     return frames
 
-def create_video_from_frames(frames, duration=5, fps=30):
+def create_video_from_frames(frames, output_filename, fps=30):
+    writer = imageio_ffmpeg.write_frames(output_filename, frames[0].shape[:2], pix_fmt_out='yuv420p', fps=fps, codec='libx264', preset="veryslow")
+    writer.send(None)
+    for frame in frames:
+        writer.send(frame)
+    writer.close()
+
+def generate_video_from_script(script, duration_per_segment=5):
+    script_segments = segment_script(script)
+    image_prompts = generate_image_prompts(script_segments)
+    all_frames = []
+
+    for i, (segment, image_prompt) in enumerate(zip(script_segments, image_prompts)):
+        print(f"Processing segment {i + 1} of {len(script_segments)}")
+        print(f"  Segment: {segment}")
+        print(f"  Image Prompt: {image_prompt}")
+        
+        image = query_hf_image_generation(image_prompt)
+        image_description = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")(image)[0]['generated_text']
+        
+        system_prompt = "You are an expert in image to video creation. Provide the motion type, intensity, text overlay, text color, text start and end times, color preset, and vignette intensity for the following image description and user prompt. Give the response in a JSON format."
+        prompt = f"Image Description: {image_description}\nUser Prompt: {segment}"
+        llm_response = query_llm(prompt, system_prompt)
+        
+        print(f"  LLM Response: {llm_response}")
+        motion_params = extract_motion_params(llm_response)
+        print(f"  Motion Parameters: {motion_params}")
+        
+        frames = apply_advanced_motion(
+            image,
+            motion_params["motion_type"],
+            motion_params["intensity"],
+            duration=duration_per_segment,
+            fps=30,
+            text_overlay=motion_params["text_overlay"],
+            text_color=motion_params["text_color"],
+            font_size=50,
+            start_time=motion_params["start_time"],
+            end_time=motion_params["end_time"],
+            color_preset=motion_params.get("color_preset", None),
+            vignette_intensity=motion_params.get("vignette_intensity", 0)
+        )
+        all_frames.extend(frames)
+
     with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmpfile:
         output_filename = tmpfile.name
-        writer = imageio_ffmpeg.write_frames(output_filename, frames[0].shape[:2], pix_fmt_out='yuv420p', fps=fps, codec='libx264', preset="veryslow")
-        writer.send(None)
-        for frame in frames:
-            writer.send(frame)
-        writer.close()
+    create_video_from_frames(all_frames, output_filename)
     return output_filename
 
 def generate_and_animate(prompt):
@@ -369,5 +433,17 @@ iface = gr.Interface(
     description="Enter a prompt to generate an image and animate it. Uses Flux 1, an LLM, and advanced video processing techniques."
 )
 
+video_iface = gr.Interface(
+    fn=generate_video_from_script,
+    inputs=[
+        gr.Textbox(label="Script (max 1 minute video)", lines=5),
+        gr.Slider(label="Duration per Segment (seconds)", minimum=1, maximum=10, step=1, value=5)
+    ],
+    outputs=gr.Video(label="Generated Video from Script"),
+    title="Story Visualizer",
+    description="Enter a short story script, and this will generate a video visualizing it using multiple images and animations."
+)
+demo = gr.TabbedInterface([iface, video_iface], ["Generate and Animate", "Story to Video"])
+
 if __name__ == "__main__":
-    iface.launch(share=True, debug=True)
+    demo.launch(share=True, debug=True)