Create app.py

app.py

@@ -0,0 +1,322 @@
+import streamlit as st
+from PIL import Image
+import io
+import google_Auth
+import mongo_db
+
+# Functions for image and prompt generation
+from image_model import (
+    flux,
+    stable_diffusion,
+)
+from prompt_model import Qwen_72b, microsoft_phi, Mixtral
+
+# Set page config with a custom title and layout
+st.set_page_config(
+    page_title="Welcome to ImagiGen: AI-Powered Image Synthesis",
+    layout="wide",
+    initial_sidebar_state="expanded",
+)
+# Custom CSS for fonts, background, and animations
+st.markdown(
+    """
+    <style>
+    /* Add a custom font */
+    @import url('https://fonts.googleapis.com/css2?family=Poppins:wght@300;400;600&display=swap');
+    body {
+        font-family: 'Poppins', sans-serif;
+        background: linear-gradient(135deg, #f8f9fa, #e9ecef);
+        animation: gradient-animation 10s ease infinite;
+    }
+
+    /* Header and title styles */
+    h1, h2, h3 {
+        color: #343a40;
+        text-align: center;
+    }
+
+    /* Button styles */
+    .stButton button {
+        background-color: #007bff;
+        color: white;
+        border-radius: 8px;
+        transition: all 0.3s ease;
+    }
+    .stButton button:hover {
+        background-color: #0056b3;
+        transform: scale(1.05);
+    }
+
+    /* Smooth gradient animation */
+    @keyframes gradient-animation {
+        0% { background: linear-gradient(135deg, #f8f9fa, #e9ecef); }
+        50% { background: linear-gradient(135deg, #e9ecef, #f8f9fa); }
+        100% { background: linear-gradient(135deg, #f8f9fa, #e9ecef); }
+    }
+
+    /* Custom input box styling */
+    input, textarea {
+        border: 2px solid #ced4da;
+        border-radius: 5px;
+        padding: 10px;
+    }
+
+    /* Sidebar customization */
+    .sidebar .sidebar-content {
+        background-color: #ffffff;
+        border-radius: 10px;
+        padding: 20px;
+    }
+    </style>
+    """,
+    unsafe_allow_html=True,
+)
+
+# List of functions
+image_functions = [flux, stable_diffusion]
+prompt_functions = [Qwen_72b, microsoft_phi, Mixtral]
+
+# Initialize session state for login status
+if "logged_in" not in st.session_state:
+    st.session_state.logged_in = False
+
+
+def register():
+    col1, col2, col3 = st.columns([1, 2, 1])  # Left, Center (main content), Right
+    with col2:
+        with st.expander("Register New Account", expanded=False):
+            st.subheader("Create a New Account", divider="rainbow")
+            new_email_id = st.text_input("email_id", placeholder="Enter an Email ID")
+            new_password = st.text_input(
+                "New Password", type="password", placeholder="Choose a password"
+            )
+            confirm_password = st.text_input(
+                "Confirm Password",
+                type="password",
+                placeholder="Re-enter your password",
+            )
+
+            if st.button("Register"):
+                if new_password != confirm_password:
+                    st.error("Passwords do not match. Please try again.")
+                elif not new_email_id or not new_password:
+                    st.error("All fields are required.")
+                else:
+                    response = mongo_db.register(new_email_id, new_password)
+                    print(response)
+                    if "Email ID already Registered" == response:
+                        st.error(
+                            "Email ID already Registered. Please choose a different email id."
+                        )
+                    elif "Registration successful" == response:
+                        st.success("Registration successful! You can now log in.")
+                        st.balloons()
+                    else:
+                        st.error("Unexpected error occur. Please Retry....")
+
+            if st.button("Register Using Google Account"):
+                user_info = google_Auth.auth()
+                user_name = f'{user_info.get("user").get("displayName")}'
+                email = f'{user_info.get("user").get("emailAddress")}'
+                response = mongo_db.google_register(user_name, email)
+                if "email ID already Registered" == response:
+                    st.error("Email ID already Registered. Please Log in.")
+                elif "Registration successful" == response:
+                    st.success("Registration successful! You can now log in.")
+                    st.balloons()
+                else:
+                    st.error("Unexpected error occur. Please Retry....")
+
+
+# Function to handle login
+def login():
+    col1, col2, col3 = st.columns([1, 2, 1])  # Left, Center (main content), Right
+    with col2:
+        _, centre, _ = st.columns([1, 2, 1])
+        with centre:
+            st.image("ImagiGen--Logo.svg", use_container_width=True)
+        st.header("Login Page", divider="rainbow")
+        st.subheader("Please Login to Continue")
+        email_id = st.text_input("Email Id", placeholder="Enter your Email Id")
+        password = st.text_input(
+            "Password", type="password", placeholder="Enter your password"
+        )
+        if st.button("Login"):
+            response = mongo_db.login(email_id, password)
+            if response == "Login successful":
+                st.session_state.logged_in = True
+                st.success("Login successful! Redirecting...")
+                st.rerun()
+            else:
+                st.error("Invalid email id or password")
+        if st.button("google"):
+            user_info = google_Auth.auth()
+            email = f'{user_info.get("user").get("emailAddress")}'
+            response = mongo_db.google_login(email)
+            if response == "Login successful":
+                st.session_state.logged_in = True
+                st.rerun()
+            else:
+                st.error("Email id is not registered. Please Register")
+
+
+# Function for the platform/dashboard page
+def platform_page():
+    _, centre, _ = st.columns([3, 4, 3])
+    with centre:
+        st.image("ImagiGen--Logo.svg", use_container_width=True)
+    st.header("Welcome to ImagiGen: AI-Powered Image Synthesis", divider="rainbow")
+    if st.button("Logout"):
+        st.session_state.logged_in = False
+        st.rerun()
+
+    st.markdown("### Describe Your Creative Vision")
+
+    # Initialize session state for prompts
+    if "model_1" not in st.session_state:
+        st.session_state.model_1 = ""
+    if "model_2" not in st.session_state:
+        st.session_state.model_2 = ""
+    if "model_3" not in st.session_state:
+        st.session_state.model_3 = ""
+    # User input for the prompt
+    Raw_user_input = st.text_input(
+        "Enter your description for prompt generation:",
+        "",
+        placeholder="E.g., A vivid image of a flowing red silk dress in a windy desert...",
+    )
+
+    user_input = f"""Instruction: Generate a highly detailed and vivid prompt based on the user input, specifically for the art and design industry, with a focus on clothing... User Input: {Raw_user_input}"""
+
+    if st.button("Generate Prompt"):
+        with st.spinner("Generating Prompts..."):
+            st.session_state.model_1 = Qwen_72b(user_input)
+            st.session_state.model_2 = microsoft_phi(user_input)
+            st.session_state.model_3 = Mixtral(user_input)
+
+    # Display text areas to edit each generated prompt
+    cols = st.columns(3)
+    with cols[0]:
+        st.session_state.model_1 = st.text_area(
+            "Edit Prompt 1:", st.session_state.model_1, height=200
+        )
+    with cols[1]:
+        st.session_state.model_2 = st.text_area(
+            "Edit Prompt 2:", st.session_state.model_2, height=200
+        )
+    with cols[2]:
+        st.session_state.model_3 = st.text_area(
+            "Edit Prompt 3:", st.session_state.model_3, height=200
+        )
+
+    if st.button("Generate Images"):
+        prompts = [
+            st.session_state.model_1,
+            st.session_state.model_2,
+            st.session_state.model_3,
+        ]
+        cols = st.columns(3)
+        with cols[0]:
+            with st.spinner(f"Generating Image {1}..."):
+                try:
+                    image_bytes = image_functions[0](prompts[0])
+                    image = Image.open(io.BytesIO(image_bytes))
+                    st.image(
+                        image,
+                        caption=f"Prompt {1}, Model {1}",
+                        use_container_width=True,
+                    )
+                except:
+                    st.image(
+                        "error.jpg",
+                        caption="Server error occur",
+                    )
+        with cols[1]:
+            with st.spinner(f"Generating Image {2}..."):
+                try:
+                    image_bytes = image_functions[0](prompts[1])
+                    image = Image.open(io.BytesIO(image_bytes))
+                    st.image(
+                        image,
+                        caption=f"Prompt {2}, Model {1}",
+                        use_container_width=True,
+                    )
+                except:
+                    st.image(
+                        "error.jpg",
+                        caption="Server error occur",
+                    )
+        with cols[2]:
+            with st.spinner(f"Generating Image {3}..."):
+                try:
+                    image_bytes = image_functions[0](prompts[2])
+                    image = Image.open(io.BytesIO(image_bytes))
+                    st.image(
+                        image,
+                        caption=f"Prompt {3}, Model {1}",
+                        use_container_width=True,
+                    )
+                except:
+                    st.image(
+                        "error.jpg",
+                        caption="Server error occur",
+                    )
+        with cols[0]:
+            with st.spinner(f"Generating Image {4}..."):
+                try:
+                    image_bytes = image_functions[1](prompts[0])
+                    image = Image.open(io.BytesIO(image_bytes))
+                    st.image(
+                        image,
+                        caption=f"Prompt {1}, Model {2}",
+                        use_container_width=True,
+                    )
+                except:
+                    st.image(
+                        "error.jpg",
+                        caption="Server error occur",
+                    )
+        with cols[1]:
+            with st.spinner(f"Generating Image {5}..."):
+                try:
+                    image_bytes = image_functions[1](prompts[2])
+                    image = Image.open(io.BytesIO(image_bytes))
+                    st.image(
+                        image,
+                        caption=f"Prompt {2}, Model {2}",
+                        use_container_width=True,
+                    )
+                except:
+                    st.image(
+                        "error.jpg",
+                        caption="Server error occur",
+                    )
+        with cols[2]:
+            with st.spinner(f"Generating Image {6}..."):
+                try:
+                    image_bytes = image_functions[1](prompts[3])
+                    image = Image.open(io.BytesIO(image_bytes))
+                    st.image(
+                        image,
+                        caption=f"Prompt {3}, Model {2}",
+                        use_container_width=True,
+                    )
+                except:
+                    st.image(
+                        "error.jpg",
+                        caption="Server error occur",
+                    )
+
+    if st.button("Clear/Regenerate"):
+        st.session_state.model_1 = ""
+        st.session_state.model_2 = ""
+        st.session_state.model_3 = ""
+        st.rerun()
+
+
+# Main logic to control which page to display
+if st.session_state.logged_in:
+    platform_page()
+else:
+    login()
+    register()