Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| from PIL import Image | |
| import io | |
| import google_auth | |
| import mongo_db | |
| # Functions for image and prompt generation | |
| from image_models import ( | |
| flux, | |
| stable_diffusion, | |
| ) | |
| from prompt_models 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. The prompt should thoroughly describe visual elements such as textures, colors, patterns, and fabric types, tailored to the context of fashion design. Incorporate details about the style, setting, and mood to effectively showcase the clothing item. The prompt should guide the creation of a realistic, lifelike image of the dress, ensuring that it is depicted from head to toe and is complete in every aspect. Emphasize that the dress must appear realistic and visually captivating. | |
| Make sure the prompt is rich in detail to enhance the image model's ability to generate an accurate and visually compelling representation. Do not include any instructions or additional details; provide only the refined prompt. | |
| 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[1]) | |
| 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[2]) | |
| 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() | |