Create app.py

app.py

@@ -0,0 +1,223 @@
+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()