Update app.py

app.py

@@ -129,13 +129,17 @@ import torch
 from langchain.chains import LLMChain
 from langchain.prompts import ChatPromptTemplate
 from langchain_huggingface import HuggingFaceEndpoint
-from diffusers import DiffusionPipeline
+
+# Additional imports for AnimateDiff
+from diffusers import AnimateDiffPipeline, MotionAdapter, EulerDiscreteScheduler
+from diffusers.utils import export_to_gif
+from huggingface_hub import hf_hub_download
+from safetensors.torch import load_file
 
 def create_conversation_prompt(name1: str, name2: str, persona_style: str):
     """
     Create a prompt that instructs the model to produce exactly 15 messages
     of conversation, alternating between name1 and name2, starting with name1.
-    We will be very explicit and not allow any formatting except the required lines.
     """
     prompt_template_str = f"""
     You are simulating a conversation of exactly 15 messages between two people: {name1} and {name2}.
@@ -171,7 +175,7 @@ def create_summary_prompt(name1: str, name2: str, conversation: str):
     return ChatPromptTemplate.from_template(summary_prompt_str)
 
 def main():
-    st.title("LLM Conversation Simulation")
+    st.title("LLM Conversation Simulation + AnimateDiff Video")
 
     model_names = [
         "meta-llama/Llama-3.3-70B-Instruct",
@@ -234,40 +238,55 @@ def main():
                 st.write(summary)
                 print("Summary:\n", summary)
 
-                # Now use the summary as a prompt for the text-to-image model
-                st.write("**Generating image from summary using ByteDance/AnimateDiff-Lightning...**")
-                print("Generating image from summary using ByteDance/AnimateDiff-Lightning...")
-
-                # Load the text-to-image pipeline
-                # Adjust this code as needed depending on GPU/CPU availability and model requirements.
-                device = "cuda" if torch.cuda.is_available() else "cpu"
-                pipe = DiffusionPipeline.from_pretrained(
-                    "ByteDance/AnimateDiff-Lightning",
-                    torch_dtype=torch.float16 if "cuda" in device else torch.float32
-                ).to(device)
-
-                # The summary prompt might contain both Title and Summary lines. Let's just use the Summary line as the prompt.
-                # Extract the Summary line
-                # Format: 
-                # Title: ...
-                # Summary: ...
-                # We'll parse the summary text to get the summary line alone
+                # Extract the summary line from the summary text
                 lines = summary.split("\n")
                 summary_line = ""
                 for line in lines:
                     if line.strip().lower().startswith("summary:"):
-                        # Extract the text after "Summary:"
                         summary_line = line.split("Summary:", 1)[-1].strip()
                         break
-
                 if not summary_line:
                     summary_line = "A friendly scene reflecting the conversation."
 
-                # Generate image using the summary line as prompt
-                image = pipe(summary_line).images[0]
+                # Now integrate AnimateDiff for text-to-video generation
+                st.write("**Generating animation from summary using ByteDance/AnimateDiff-Lightning...**")
+                print("Generating animation from summary using ByteDance/AnimateDiff-Lightning...")
+
+                device = "cuda" if torch.cuda.is_available() else "cpu"
+                dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+
+                step = 4  # Adjust if needed
+                repo = "ByteDance/AnimateDiff-Lightning"
+                ckpt = f"animatediff_lightning_{step}step_diffusers.safetensors"
+                base = "emilianJR/epiCRealism"  # Check if this model exists or choose a known base model
+
+                # Load and configure AnimateDiff pipeline
+                adapter = MotionAdapter().to(device, dtype)
+                adapter.load_state_dict(load_file(hf_hub_download(repo ,ckpt), device=device))
+
+                pipe = AnimateDiffPipeline.from_pretrained(
+                    base,
+                    motion_adapter=adapter,
+                    torch_dtype=dtype
+                ).to(device)
+
+                pipe.scheduler = EulerDiscreteScheduler.from_config(
+                    pipe.scheduler.config,
+                    timestep_spacing="trailing",
+                    beta_schedule="linear"
+                )
+
+                # Generate the animation
+                output = pipe(prompt=summary_line, guidance_scale=1.0, num_inference_steps=step)
+
+                # Save as GIF
+                # output.frames is a list of frames (PIL images)
+                st.write("**Exporting animation to GIF...**")
+                print("Exporting animation to GIF...")
+                export_to_gif(output.frames, "animation.gif")
 
-                st.subheader("Generated Image:")
-                st.image(image, caption="Generated by ByteDance/AnimateDiff-Lightning")
+                st.subheader("Generated Animation:")
+                st.image("animation.gif", caption="Generated by AnimateDiff using summary prompt")
 
             except Exception as e:
                 st.error(f"Error generating conversation or summary: {e}")