app.py

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()