Update app.py

app.py

@@ -1,223 +1,41 @@
 import streamlit as st
-import firebase_admin
-from firebase_admin import credentials
-from firebase_admin import auth
-from firebase_admin import firestore
-from firebase_admin import storage
-import requests
-from io import BytesIO
-from PIL import Image
-import os
-import tempfile
-import mimetypes
-import uuid
-
-
-TOKEN = os.getenv("TOKEN0")
-API_URL = os.getenv("API_URL")
-token_id = 0
-tokens_tried = 0
-no_of_accounts = 11
-model_id = os.getenv("MODEL_ID")
-
-# Initialize Firebase Admin SDK
-cred = credentials.Certificate("path/to/your/firebase/credentials.json")
-firebase_admin.initialize_app(cred, {
-    'storageBucket': 'your_storage_bucket_url'
-})
-db = firestore.client()
-bucket = storage.bucket()
-
-
-def get_image_from_url(url):
-    """
-    Fetches and returns an image from a given URL, converting to PNG if needed.
-    """
-    try:
-        response = requests.get(url, stream=True)
-        response.raise_for_status()
-        image = Image.open(BytesIO(response.content))
-        return image, url  # Return the image and the URL
-
-    except requests.exceptions.RequestException as e:
-        return f"Error fetching image: {e}", None
-    except Exception as e:
-        return f"Error processing image: {e}", None
-
-
-def generate_image(prompt, aspect_ratio, realism, user_id):
-    global token_id
-    global TOKEN
-    global tokens_tried
-    global no_of_accounts
-    global model_id
-
-    payload = {
-        "id": model_id,
-        "inputs": [prompt, aspect_ratio, str(realism).lower()],
-    }
-    headers = {"Authorization": f"Bearer {TOKEN}"}
-
-    try:
-        response_data = requests.post(API_URL, json=payload, headers=headers).json()
-        if "error" in response_data:
-            if 'error 429' in response_data['error']:
-                if tokens_tried < no_of_accounts:
-                    token_id = (token_id + 1) % (no_of_accounts)
-                    tokens_tried += 1
-                    TOKEN = os.getenv(f"TOKEN{token_id}")
-                    response_data = generate_image(prompt, aspect_ratio, realism, user_id)
-                    tokens_tried = 0
-                    return response_data
-                return "No credits available", None, None
-            return response_data, None, None
-        elif "output" in response_data:
-            url = response_data['output']
-            image, url = get_image_from_url(url)
-            # Store in firebase
-            download_path = download_image(url)
-            if download_path:
-                try:
-                    # Create a unique filename for the image
-                    image_filename = f"{uuid.uuid4()}{os.path.splitext(download_path)[1]}"
-                    # Store file in cloud storage
-                    blob = bucket.blob(f'users/{user_id}/{image_filename}')
-                    blob.upload_from_filename(download_path)
-                    image_url = blob.public_url
-                    # Store prompt and image URL in Firestore
-                    doc_ref = db.collection("images").add({
-                        "user_id": user_id,
-                        "prompt": prompt,
-                        "image_url": image_url,
-                    })
-                    print(f"Document added with id: {doc_ref.id}")
-                    os.remove(download_path)
-                    return image, image_url, image_url
-                except Exception as e:
-                    print(f"Error uploading image to firebase {e}")
-                    return image, None, url
-            else:
-                return image, None, url  # Return the image and the URL
-        else:
-            return "Error: Unexpected response from server", None, None
-    except Exception as e:
-        print(f"Error: {e}")
-        return f"Error", None, None
-
-
-def download_image(image_url):
-    if not image_url:
-        return None  # Return None if image_url is empty
-    try:
-        response = requests.get(image_url, stream=True)
-        response.raise_for_status()
-
-        # Get the content type from the headers
-        content_type = response.headers.get('content-type')
-        extension = mimetypes.guess_extension(content_type)
-
-        if not extension:
-            extension = ".png"  # Default to .png if can't determine the extension
-
-        # Create a temporary file with the correct extension
-        with tempfile.NamedTemporaryFile(suffix=extension, delete=False) as tmp_file:
-            for chunk in response.iter_content(chunk_size=8192):
-                tmp_file.write(chunk)
-            temp_file_path = tmp_file.name
-        return temp_file_path
-    except Exception as e:
-        return None
-
-
-def create_user_with_email_and_password(email, password):
-    try:
-       user = auth.create_user(
-          email = email,
-          password = password,
-      )
-       return user.uid, None
-    except Exception as e:
-        print(f"Error creating user {e}")
-        return None, str(e)
-
-def sign_in_with_email_and_password(email, password):
-     try:
-         user = auth.get_user_by_email(email)
-         return user.uid, None
-     except Exception as e:
-        print(f"Error logging in user: {e}")
-        return None, str(e)
-
-
-def main():
-    st.title("Image Generator")
-
-    # Initialize session state
-    if "user_id" not in st.session_state:
-        st.session_state.user_id = None
-
-
-    # Check if user is logged in
-    if not st.session_state.user_id:
-          st.subheader("Login/Signup")
-          st.session_state.login_status = False
-          login_option = st.radio("Select Option", ("Login", "Sign Up"))
-          if login_option == "Login":
-            email = st.text_input("Email")
-            password = st.text_input("Password", type="password")
-            if st.button("Log In"):
-                if email and password:
-                    user_id, error = sign_in_with_email_and_password(email, password)
-                    if user_id:
-                        st.success("Logged in successfully!")
-                        st.session_state.user_id = user_id
-                        st.session_state.login_status = True
-                    else:
-                        st.error(f"Login Failed: {error}")
-                else:
-                    st.error("Please provide both email and password")
-          elif login_option == "Sign Up":
-              email = st.text_input("Email")
-              password = st.text_input("Password", type="password")
-              if st.button("Sign Up"):
-                  if email and password:
-                      user_id, error = create_user_with_email_and_password(email, password)
-                      if user_id:
-                          st.success("Signed up successfully!")
-                          st.session_state.user_id = user_id
-                          st.session_state.login_status = True
-                      else:
-                          st.error(f"Sign Up Failed: {error}")
-                  else:
-                      st.error("Please provide both email and password")
-
-    # Main app logic if logged in
-    if st.session_state.user_id:
-        prompt = st.text_input("Prompt", placeholder="Describe the image you want to generate")
-        aspect_ratio = st.radio(
-            "Aspect Ratio",
-            ["1:1", "3:4", "4:3", "9:16", "16:9", "9:21", "21:9"],
-            index=4 #Default value is 16:9
-        )
-        realism = st.checkbox("Realism", value=False)
-
-        if st.button("Generate Image"):
-            image, url, file_url = generate_image(prompt, aspect_ratio, realism, st.session_state.user_id)
-
-            if isinstance(image, str):
-               st.error(image)
-            elif image:
-                st.image(image, caption="Generated Image", use_column_width=True)
-                if url:
-                   st.write(f"Firestore URL: {url}")
-                if file_url:
-                   st.download_button(label="Download Image", data=requests.get(file_url).content, file_name="image.jpg", mime="image/jpeg")
-
-
-        if st.button("Log Out"):
-            st.session_state.user_id = None
-            st.rerun()
-
-
-if __name__ == "__main__":
-    main()
+import streamlit_authenticator as stauth
+
+# Define user authentication
+users = {
+    "credentials": {
+        "usernames": {
+            "user1": {"name": "User One", "password": "password1"},
+            "user2": {"name": "User Two", "password": "password2"},
+        }
+    },
+    "cookie": {
+        "expiry_days": 30,
+        "key": "my_secret_key",
+        "name": "streamlit_auth",
+    },
+}
+
+# Initialize the authenticator
+authenticator = stauth.Authenticate(
+    credentials=users["credentials"],
+    cookie_name=users["cookie"]["name"],
+    key=users["cookie"]["key"],
+    cookie_expiry_days=users["cookie"]["expiry_days"],
+)
+
+# Login/Logout flow
+name, authentication_status, username = authenticator.login("Login", "main")
+
+if authentication_status:
+    st.success(f"Welcome, {name}!")
+    user_input = st.text_input("Type a message here:")
+    if user_input:
+        st.write(f"Response: hi")
+    authenticator.logout("Logout", "sidebar")
+elif authentication_status == False:
+    st.error("Username/password is incorrect.")
+elif authentication_status == None:
+    st.warning("Please enter your username and password.")
+
+# Note: You can replace `stauth` with your custom Google OAuth implementation for more complex use.