Update app.py

app.py

@@ -1,61 +1,99 @@
-from transformers import AutoTokenizer
-import transformers
 import torch
 import streamlit as st
+from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
+from diffusers import StableDiffusionPipeline
+from PIL import Image
 
-model_name = "tiiuae/falcon-7b-instruct"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
+text_model = "gpt2"
+tokenizer = AutoTokenizer.from_pretrained(text_model)
+model = AutoModelForCausalLM.from_pretrained(text_model)
 
-generator = transformers.pipeline(
+image_model = "runwayml/stable-diffusion-v1-5"
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+generator = pipeline(
     "text-generation",
-    model=model_name,
+    model=model,
     tokenizer=tokenizer,
-    torch_dtype=torch.bfloat16,
-    trust_remote_code=True,
-    device_map="auto"
+    device=0 if torch.cuda.is_available() else -1
 )
 
-def generate_text(prompt, section, max_length=200, top_k=50, temperature=0.7):
-    sequences = generator(
-        f"{section} - {prompt}",
+def generate_text(prompt, temperature=0.7, top_k=50, repetition_penalty=1.2, max_length=None, min_length=10):
+    return generator(
+        prompt,
         max_length=max_length,
+        min_length=min_length,
         do_sample=True,
         top_k=top_k,
         temperature=temperature,
+        repetition_penalty=repetition_penalty,
         num_return_sequences=1,
         eos_token_id=tokenizer.eos_token_id,
-    )
-    return sequences[0]["generated_text"]
+    )[0]["generated_text"]
+
+def generate_image(prompt):
+    pipe = StableDiffusionPipeline.from_pretrained(image_model, torch_dtype=torch.float32)
+    pipe = pipe.to(device)
+    image = pipe(prompt).images[0]
+    return image
 
-# Streamlit app
 st.title("AI-Generated Blog Post")
 
-# Keyword selection input
-keywords_input = st.text_input("Step 1: Keyword Selection (Separate keywords with commas)", "Artificial Intelligence")
+title = st.text_input("Topic of the Article")
+keywords_input = st.text_input("Enter Some Keywords About The Topic (Separate keywords with commas)")
 keywords = [word.strip() for word in keywords_input.split(',')]
+keywords.append(title)
 
-# Display generated content on button click
 if st.button('Generate Article'):
-    if keywords_input:
-
+    if keywords:
         generated_text = " ".join(keywords)
-        intro_text = generate_text(generated_text, "Introduction", max_length=200, top_k=50, temperature=0.7)
-        body_text = generate_text(generated_text, "Body", max_length=500, top_k=50, temperature=0.7)
-        conclusion_text = generate_text(generated_text, "Conclusion", max_length=150, top_k=50, temperature=0.7)
 
-        # Displaying the sections
+        formatted_title = title.capitalize()
+
+        st.markdown(
+            f"<h1 style='text-align: center; color: blue; font-size: 70px;'>{formatted_title}</h1>",
+            unsafe_allow_html=True
+        )
+
+        generated_image1 = generate_image(generated_text)
+
+        col1, col2, col3 = st.columns([1, 2, 1])
+        with col2:
+            new_image1 = generated_image1.resize((700, 300))  # Resize the image here
+            st.image(new_image1, use_column_width=True)
+
+        intro_text = generate_text(generated_text, min_length=100, max_length=200)
+        body_text1 = generate_text(generated_text, min_length=200, max_length=250)
+        body_text2 = generate_text(generated_text, min_length=300, max_length=400)
+        conclusion_text = generate_text(generated_text, min_length=100, max_length=200)
+
         st.header("Introduction")
         st.write(intro_text)
+        modified_prompt = generated_text + ' bright'
+        generated_image2 = generate_image(modified_prompt)
+
+        new_image2 = generated_image1.resize((700, 300))  # Resize the image here
+        st.image(new_image2, use_column_width=True)
+
+
+        col1, col2 = st.columns(2)
+        with col1:
+            st.header("Body")
+            st.write(body_text1)
+
+        with col2:
+            modified_prompt2 = generated_text + ' shade'
+            generated_image3 = generate_image(modified_prompt2)
+            st.markdown("<br><br><br><br>", unsafe_allow_html=True)  # Add vertical space
+            st.image(generated_image3, use_column_width=True)
 
-        st.header("Body")
-        st.write(body_text)
+        st.write(body_text2)
 
         st.header("Conclusion")
         st.write(conclusion_text)
     else:
         st.warning("Please input keywords to generate content.")
 
-# Sidebar with instructions
 st.sidebar.title("Instructions")
 st.sidebar.write(
     "1. Enter keywords related to the topic you want to generate content about."