Update app.py

app.py

@@ -1,22 +1,146 @@
+# import os
+# import streamlit as st
+# import torch
+# from langchain.chains import LLMChain
+# from langchain.prompts import ChatPromptTemplate
+# from langchain_huggingface import HuggingFaceEndpoint
+
+# 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}.
+#     {name1} speaks first (message 1), then {name2} (message 2), then {name1} (message 3), and so forth,
+#     alternating until all 15 messages are complete. The 15th message is by {name1}.
+
+#     Requirements:
+#     - Output exactly 15 lines, no more, no less.
+#     - Each line must be a single message in the format:
+#       {name1}: <message> or {name2}: <message>
+#     - Do not add any headings, numbers, sample outputs, or explanations.
+#     - Do not mention code, programming, or instructions.
+#     - Each message should be 1-2 short sentences, friendly, natural, reflecting the style: {persona_style}.
+#     - Use everyday language, can ask questions, show opinions.
+#     - Use emojis sparingly if it fits the style (no more than 1-2 total).
+#     - No repeated lines, each message should logically follow from the previous one.
+#     - Do not produce anything after the 15th message. No extra lines or text.
+
+#     Produce all 15 messages now:
+#     """
+#     return ChatPromptTemplate.from_template(prompt_template_str)
+
+# def create_summary_prompt(name1: str, name2: str, conversation: str):
+#     """Prompt for generating a title and summary."""
+#     summary_prompt_str = f"""
+#     Below is a completed 15-message conversation between {name1} and {name2}:
+
+#     {conversation}
+
+#     Please provide:
+#     Title: <A short descriptive title of the conversation>
+#     Summary: <A few short sentences highlighting the main points, tone, and conclusion>
+
+#     Do not continue the conversation, do not repeat it, and do not add extra formatting beyond the two lines:
+#     - One line starting with "Title:"
+#     - One line starting with "Summary:"
+#     """
+#     return ChatPromptTemplate.from_template(summary_prompt_str)
+
+# def main():
+#     st.title("LLM Conversation Simulation")
+
+#     model_names = [
+#         "meta-llama/Llama-3.3-70B-Instruct",
+#         "mistralai/Mistral-7B-v0.1",
+#         "tiiuae/falcon-7b"
+        
+#     ]
+#     selected_model = st.selectbox("Select a model:", model_names)
+
+#     name1 = st.text_input("Enter the first user's name:", value="Alice")
+#     name2 = st.text_input("Enter the second user's name:", value="Bob")
+#     persona_style = st.text_area("Enter the persona style characteristics:", 
+#                                  value="friendly, curious, and a bit sarcastic")
+
+#     if st.button("Start Conversation Simulation"):
+#         st.write("**Loading model...**")
+#         print("Loading model...")
+
+#         with st.spinner("Starting simulation..."):
+#             endpoint_url = f"https://api-inference.huggingface.co/models/{selected_model}"
+
+#             try:
+#                 llm = HuggingFaceEndpoint(
+#                     endpoint_url=endpoint_url,
+#                     huggingfacehub_api_token=os.environ.get("HUGGINGFACEHUB_API_TOKEN"),
+#                     task="text-generation",
+#                     temperature=0.7,
+#                     max_new_tokens=512
+#                 )
+#                 st.write("**Model loaded successfully!**")
+#                 print("Model loaded successfully!")
+#             except Exception as e:
+#                 st.error(f"Error initializing HuggingFaceEndpoint: {e}")
+#                 print(f"Error initializing HuggingFaceEndpoint: {e}")
+#                 return
+
+#             conversation_prompt = create_conversation_prompt(name1, name2, persona_style)
+#             conversation_chain = LLMChain(llm=llm, prompt=conversation_prompt)
+
+#             st.write("**Generating the full 15-message conversation...**")
+#             print("Generating the full 15-message conversation...")
+
+#             try:
+#                 # Generate all 15 messages in one go
+#                 conversation = conversation_chain.run(chat_history="", input="").strip()
+
+#                 st.subheader("Final Conversation:")
+#                 st.text(conversation)
+#                 print("Conversation Generation Complete.\n")
+#                 print("Full Conversation:\n", conversation)
+
+#                 # Summarize the conversation
+#                 summary_prompt = create_summary_prompt(name1, name2, conversation)
+#                 summary_chain = LLMChain(llm=llm, prompt=summary_prompt)
+
+#                 st.subheader("Summary and Title:")
+#                 st.write("**Summarizing the conversation...**")
+#                 print("Summarizing the conversation...")
+
+#                 summary = summary_chain.run(chat_history="", input="")
+#                 st.write(summary)
+#                 print("Summary:\n", summary)
+
+#             except Exception as e:
+#                 st.error(f"Error generating conversation: {e}")
+#                 print(f"Error generating conversation: {e}")
+
+# if __name__ == "__main__":
+#     main()
+
+
 import os
 import streamlit as st
 import torch
 from langchain.chains import LLMChain
 from langchain.prompts import ChatPromptTemplate
 from langchain_huggingface import HuggingFaceEndpoint
+from diffusers import DiffusionPipeline
 
 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}.
     {name1} speaks first (message 1), then {name2} (message 2), then {name1} (message 3), and so forth,
     alternating until all 15 messages are complete. The 15th message is by {name1}.
-
     Requirements:
     - Output exactly 15 lines, no more, no less.
     - Each line must be a single message in the format:
@@ -28,7 +152,6 @@ def create_conversation_prompt(name1: str, name2: str, persona_style: str):
     - Use emojis sparingly if it fits the style (no more than 1-2 total).
     - No repeated lines, each message should logically follow from the previous one.
     - Do not produce anything after the 15th message. No extra lines or text.
-
     Produce all 15 messages now:
     """
     return ChatPromptTemplate.from_template(prompt_template_str)
@@ -37,13 +160,10 @@ def create_summary_prompt(name1: str, name2: str, conversation: str):
     """Prompt for generating a title and summary."""
     summary_prompt_str = f"""
     Below is a completed 15-message conversation between {name1} and {name2}:
-
     {conversation}
-
     Please provide:
     Title: <A short descriptive title of the conversation>
     Summary: <A few short sentences highlighting the main points, tone, and conclusion>
-
     Do not continue the conversation, do not repeat it, and do not add extra formatting beyond the two lines:
     - One line starting with "Title:"
     - One line starting with "Summary:"
@@ -57,7 +177,6 @@ def main():
         "meta-llama/Llama-3.3-70B-Instruct",
         "mistralai/Mistral-7B-v0.1",
         "tiiuae/falcon-7b"
-        
     ]
     selected_model = st.selectbox("Select a model:", model_names)
 
@@ -115,9 +234,45 @@ 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
+                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]
+
+                st.subheader("Generated Image:")
+                st.image(image, caption="Generated by ByteDance/AnimateDiff-Lightning")
+
             except Exception as e:
-                st.error(f"Error generating conversation: {e}")
-                print(f"Error generating conversation: {e}")
+                st.error(f"Error generating conversation or summary: {e}")
+                print(f"Error generating conversation or summary: {e}")
 
 if __name__ == "__main__":
     main()
+