Update utils.py

utils.py

@@ -1,26 +1,62 @@
-import gradio as gr
+import torch
+from diffusers import StableDiffusionPipeline
+from transformers import AutoModelForCausalLM, AutoTokenizer, AutoProcessor
+import os, random
+from moviepy.editor import ImageSequenceClip, AudioFileClip
 
-def create_story_video(prompt, duration):
-    story = generate_story(prompt, duration)
-    audio_path = generate_voice(story)
-    image_paths = generate_images(story)
-    video_path = make_video(image_paths, audio_path)
+# ---------- STORY GENERATION ----------
+def generate_story(prompt, duration):
+    model = AutoModelForCausalLM.from_pretrained(
+        "mistralai/Mistral-7B-Instruct-v0.2", torch_dtype=torch.float16, device_map="auto"
+    )
+    tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.2")
 
-    return story, audio_path, video_path
+    text = tokenizer(
+        f"Write a detailed story of {duration} minutes about: {prompt}", 
+        return_tensors="pt"
+    ).to("cuda")
 
-with gr.Blocks() as app:
-    gr.Markdown("# 🎬 AI Story Video Generator")
-    gr.Markdown("Generate 5–10 minute story videos with AI images + AI voice.")
+    output = model.generate(**text, max_new_tokens=900)
+    story = tokenizer.decode(output[0], skip_special_tokens=True)
+    return story
 
-    prompt = gr.Textbox(label="Story Topic", placeholder="ex: A lonely robot on Mars finds a friend...")
-    duration = gr.Slider(1, 10, value=5, step=1, label="Story duration (minutes)")
-    btn = gr.Button("Generate Story Video")
 
-    story_output = gr.Textbox(label="Generated Story")
-    audio_output = gr.Audio(label="AI Narration")
-    video_output = gr.Video(label="Final AI Story Video")
+# ---------- VOICE ----------
+def generate_voice(text):
+    from bark import generate_audio, preload_models
+    preload_models()
+    audio = generate_audio(text)
+    path = "story_audio.wav"
+    import soundfile as sf
+    sf.write(path, audio, 24000)
+    return path
 
-    btn.click(create_story_video, inputs=[prompt, duration],
-              outputs=[story_output, audio_output, video_output])
 
-app.launch()
+# ---------- IMAGE GENERATION ----------
+def generate_images(story):
+    sd = StableDiffusionPipeline.from_pretrained(
+        "runwayml/stable-diffusion-v1-5",
+        torch_dtype=torch.float16
+    ).to("cuda")
+
+    scenes = story.split(".")[:10]   # create 10 images max
+    paths = []
+
+    for i, s in enumerate(scenes):
+        img = sd(s).images[0]
+        path = f"img_{i}.png"
+        img.save(path)
+        paths.append(path)
+
+    return paths
+
+
+# ---------- VIDEO ----------
+def make_video(images, audio_file):
+    video_clip = ImageSequenceClip(images, fps=1)  # 1 image per second
+    audio_clip = AudioFileClip(audio_file)
+    final = video_clip.set_audio(audio_clip)
+
+    out = "final_story_video.mp4"
+    final.write_videofile(out)
+    return out