Update app.py

app.py

@@ -1,13 +1,20 @@
 import streamlit as st
 import firebase_admin
-from firebase_admin import credentials, auth, db
+from firebase_admin import credentials, auth, db, storage
 import os
 import json
 import requests
+from io import BytesIO
+from PIL import Image
+import tempfile
+import mimetypes
+import uuid
 
 # Load Firebase credentials from Hugging Face Secrets
 firebase_creds = os.getenv("FIREBASE_CREDENTIALS")
 FIREBASE_API_KEY = os.getenv("FIREBASE_API_KEY")
+FIREBASE_STORAGE_BUCKET = os.getenv("FIREBASE_STORAGE_BUCKET")
+
 if firebase_creds:
     firebase_creds = json.loads(firebase_creds)
 else:
@@ -17,7 +24,8 @@ else:
 if not firebase_admin._apps:
     cred = credentials.Certificate(firebase_creds)
     firebase_admin.initialize_app(cred, {
-        'databaseURL': 'https://creative-623ef-default-rtdb.firebaseio.com/'
+        'databaseURL': 'https://creative-623ef-default-rtdb.firebaseio.com/',
+        'storageBucket': FIREBASE_STORAGE_BUCKET
     })
 
 # Initialize session state
@@ -28,6 +36,12 @@ if "current_user" not in st.session_state:
 if "display_name" not in st.session_state:
     st.session_state.display_name = None
 
+TOKEN = os.getenv("TOKEN0")
+API_URL = os.getenv("API_URL")
+token_id = 0
+tokens_tried = 0
+no_of_accounts = 7
+model_id = os.getenv("MODEL_ID")
 
 def send_verification_email(id_token):
     url = f'https://identitytoolkit.googleapis.com/v1/accounts:sendOobCode?key={FIREBASE_API_KEY}'
@@ -115,18 +129,117 @@ def logout_callback():
     st.session_state.display_name = None
     st.info("Logged out successfully!")
 
-# Function to store text in Firebase Realtime Database
-def store_text_in_db(user_id, text):
+# Function to get image from url
+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
+
+# Function to generate image
+def generate_image(prompt, aspect_ratio, realism):
+    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)
+                    tokens_tried = 0
+                    return response_data
+                return "No credits available", None
+            return response_data, None
+        elif "output" in response_data:
+            url = response_data['output']
+            image, url = get_image_from_url(url)
+            return image, url  # Return the image and the URL
+        else:
+            return "Error: Unexpected response from server", None
+    except Exception as e:
+        return f"Error", 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
+
+# Function to store image and related data in Firebase
+def store_image_data_in_db(user_id, prompt, aspect_ratio, realism, image_url):
     try:
-        ref = db.reference(f'users/{user_id}/texts')
-        new_text_ref = ref.push()
-        new_text_ref.set({
-            'text': text,
+        ref = db.reference(f'users/{user_id}/images')
+        new_image_ref = ref.push()
+        new_image_ref.set({
+            'prompt': prompt,
+            'aspect_ratio': aspect_ratio,
+            'realism': realism,
+            'image_url': image_url,
             'timestamp': {'.sv': 'timestamp'}
         })
-        st.success("Text saved successfully!")
+        st.success("Image data saved successfully!")
     except Exception as e:
-        st.error(f"Failed to save text: {e}")
+        st.error(f"Failed to save image data: {e}")
+
+#Function to upload image to cloud storage
+def upload_image_to_storage(image, user_id):
+        try:
+            bucket = storage.bucket()
+            image_id = str(uuid.uuid4())
+            file_path = f"user_images/{user_id}/{image_id}.png"
+            blob = bucket.blob(file_path)
+
+            # Convert PIL Image to BytesIO object
+            img_byte_arr = BytesIO()
+            image.save(img_byte_arr, format='PNG')
+            img_byte_arr = img_byte_arr.getvalue()
+
+            blob.upload_from_string(img_byte_arr, content_type='image/png')
+            blob.make_public()
+            image_url = blob.public_url
+            return image_url
+        except Exception as e:
+            st.error(f"Failed to upload image to cloud storage: {e}")
+            return None
 
 # Registration form
 def registration_form():
@@ -148,16 +261,40 @@ def login_form():
 # Main app screen (after login)
 def main_app():
     st.subheader(f"Welcome, {st.session_state.display_name}!")
-    st.write("Enter a name below to get a greeting.")
+    st.write("Enter a prompt below to generate an image.")
+
+    # Input fields
+    prompt = st.text_input("Prompt", key="image_prompt", placeholder="Describe the image you want to generate")
+    aspect_ratio = st.radio(
+        "Aspect Ratio", 
+        options=["1:1", "3:4", "4:3", "9:16", "16:9", "9:21", "21:9"], 
+        index=5
+    )
+    realism = st.checkbox("Realism", value=False)
 
-    # Input field for name
-    name = st.text_input("Enter a name") # Removed key here
+    if st.button("Generate Image"):
+        if prompt:
+            with st.spinner("Generating Image..."):
+                image_result = generate_image(prompt, aspect_ratio, realism)
+                if isinstance(image_result, tuple) and len(image_result) == 2 and image_result[0] is not None:
+                    image, image_url = image_result
+                    st.image(image, caption="Generated Image", use_column_width=True)
 
-    # Display greeting in real-time
-    if name:
-        st.write(f"Hello {name}!")
-        # Store the entered text in Firebase Realtime Database
-        store_text_in_db(st.session_state.current_user, name)
+                     # Upload image to cloud storage and store url
+                    cloud_storage_url = upload_image_to_storage(image, st.session_state.current_user)
+
+                    if cloud_storage_url:
+                        # Store image data in database
+                        store_image_data_in_db(st.session_state.current_user, prompt, aspect_ratio, realism, cloud_storage_url)
+                        st.success("Image stored to database successfully!")
+
+                    download_path = download_image(image_url)
+                    if download_path:
+                        st.download_button(label="Download Image", data = open(download_path, "rb"), file_name = f"image.png")
+                else:
+                    st.error(f"Image generation failed: {image_result}")
+        else:
+            st.warning("Please enter a prompt to generate an image.")
 
     # Logout button
     if st.button("Logout", on_click=logout_callback):