Update app.py

app.py

@@ -4,14 +4,14 @@ import random
 import sys
 import torch
 from PIL import Image, ImageDraw, ImageFont
-from diffusers import DiffusionPipeline, DPMSolverMultistepScheduler
+from diffusers import DiffusionPipeline, TCDScheduler
 from huggingface_hub import hf_hub_download
 from gtts import gTTS
-from moviepy.editor import ImageSequenceClip, concatenate_videoclips, AudioFileClip
+from moviepy.editor import ImageSequenceClip, VideoFileClip, concatenate_videoclips, AudioFileClip
 import gradio as gr
 
-# Set the number of inference steps to a lower value
-num_inference_steps = 25
+# Choose among 1, 2, 4 and 8:
+num_inference_steps = 8
 
 base_model_id = "stabilityai/stable-diffusion-xl-base-1.0"
 repo_name = "ByteDance/Hyper-SD"
@@ -22,16 +22,18 @@ device = "cpu"
 pipe = DiffusionPipeline.from_pretrained(base_model_id).to(device)
 pipe.load_lora_weights(hf_hub_download(repo_name, ckpt_name))
 pipe.fuse_lora()
-pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
+pipe.scheduler = TCDScheduler.from_config(pipe.scheduler.config)
 
-def generate_image(prompt, step_count=num_inference_steps, seed=None):
+def generate_image(prompt, step_count=50, seed=None):
     if seed is None:
         seed = random.randint(0, sys.maxsize)
     generator = torch.Generator(device).manual_seed(seed)
+    eta = 0.5
     images = pipe(
         prompt=prompt,
         num_inference_steps=step_count,
-        guidance_scale=7.5,  # Increase guidance scale to improve quality with fewer steps
+        guidance_scale=0.0,
+        eta=eta,
         generator=generator,
     ).images
     return images[0]
@@ -86,7 +88,7 @@ def process_story(story):
     for i, sentence in enumerate(sentences):
         print(f"Sentence {i+1}: {sentence}\n")
         seed = random.randint(0, sys.maxsize)
-        image = generate_image(sentence, step_count=num_inference_steps, seed=seed)  # Adjusted step count
+        image = generate_image(sentence, step_count=50, seed=seed)  # Increase step count for better quality images
         resized_image = image.resize((256, 256))
         image_with_text = draw_text_on_image(resized_image, sentence)
 
@@ -113,6 +115,9 @@ def process_story(story):
 
         video_clips.append(video_clip)
 
+        # Clear memory
+        del resized_image, image_with_text
+
     # Concatenate all video clips into a final video
     final_video = concatenate_videoclips(video_clips)
     final_video_path = "story_video.mp4"