artintel235/creative2
0
1import streamlit as st2import firebase_admin3from firebase_admin import credentials, auth, db, storage4import os5import json6import requests7from io import BytesIO8from PIL import Image9import tempfile10import mimetypes11import uuid12import io13 14# Load Firebase credentials from Hugging Face Secrets15firebase_creds = os.getenv("FIREBASE_CREDENTIALS")16FIREBASE_API_KEY = os.getenv("FIREBASE_API_KEY")17FIREBASE_STORAGE_BUCKET = os.getenv("FIREBASE_STORAGE_BUCKET")18 19if firebase_creds:20 firebase_creds = json.loads(firebase_creds)21else:22 st.error("Firebase credentials not found. Please check your secrets.")23 24# Initialize Firebase (only once)25if not firebase_admin._apps:26 cred = credentials.Certificate(firebase_creds)27 firebase_admin.initialize_app(cred, {28 'databaseURL': 'https://creative-623ef-default-rtdb.firebaseio.com/',29 'storageBucket': FIREBASE_STORAGE_BUCKET30 })31 32# Initialize session state33if "logged_in" not in st.session_state:34 st.session_state.logged_in = False35if "current_user" not in st.session_state:36 st.session_state.current_user = None37if "display_name" not in st.session_state:38 st.session_state.display_name = None39if "window_size" not in st.session_state:40 st.session_state.window_size = 541if "current_window_start" not in st.session_state:42 st.session_state.current_window_start = 043if "selected_image" not in st.session_state:44 st.session_state.selected_image = None45 46TOKEN = os.getenv("TOKEN0")47API_URL = os.getenv("API_URL")48token_id = 049tokens_tried = 050no_of_accounts = 751model_id = os.getenv("MODEL_ID")52 53def send_verification_email(id_token):54 url = f'https://identitytoolkit.googleapis.com/v1/accounts:sendOobCode?key={FIREBASE_API_KEY}'55 headers = {'Content-Type': 'application/json'}56 data = {57 'requestType': 'VERIFY_EMAIL',58 'idToken': id_token59 }60 61 response = requests.post(url, headers=headers, json=data)62 result = response.json()63 64 if 'error' in result:65 return {'status': 'error', 'message': result['error']['message']}66 else:67 return {'status': 'success', 'email': result['email']}68 69# Callback for registration70def register_callback():71 email = st.session_state.reg_email72 password = st.session_state.reg_password73 display_name = st.session_state.reg_display_name74 try:75 # Step 1: Create a new user in Firebase76 user = auth.create_user(email=email, password=password)77 78 # Step 2: Update the user profile with the display name79 auth.update_user(user.uid, display_name=display_name)80 81 st.success("Registration successful! Sending verification email...")82 83 # Step 3: Sign in the user programmatically to get the id_token84 url = f'https://identitytoolkit.googleapis.com/v1/accounts:signInWithPassword?key={FIREBASE_API_KEY}'85 data = {86 'email': email,87 'password': password,88 'returnSecureToken': True89 }90 response = requests.post(url, json=data)91 result = response.json()92 93 if 'idToken' in result:94 id_token = result['idToken']95 st.session_state.id_token = id_token96 97 verification_result = send_verification_email(id_token)98 if verification_result['status'] == 'success':99 st.success(f"Verification email sent to {email}.")100 else:101 st.error(f"Failed to send verification email: {verification_result['message']}")102 else:103 st.error(f"Failed to retrieve id_token: {result['error']['message']}")104 except Exception as e:105 st.error(f"Registration failed: {e}")106 107# Callback for login108def login_callback():109 login_identifier = st.session_state.login_identifier110 password = st.session_state.login_password111 try:112 # Try to sign in the user programmatically to check the password validity113 url = f'https://identitytoolkit.googleapis.com/v1/accounts:signInWithPassword?key={FIREBASE_API_KEY}'114 data = {115 'email': login_identifier,116 'password': password,117 'returnSecureToken': True118 }119 response = requests.post(url, json=data)120 result = response.json()121 122 if 'idToken' in result:123 # If sign in was successful, then use email to fetch the user124 user = auth.get_user_by_email(login_identifier)125 st.session_state.logged_in = True126 st.session_state.current_user = user.uid127 st.session_state.display_name = user.display_name # Store the display name128 st.success("Logged in successfully!")129 130 elif 'error' in result:131 # If sign-in fails, retrieve user using display name132 try:133 user_list = auth.list_users()134 for user_info in user_list.users:135 if user_info.display_name == login_identifier:136 user = user_info137 # If user is found using display name, try signing in using email138 url = f'https://identitytoolkit.googleapis.com/v1/accounts:signInWithPassword?key={FIREBASE_API_KEY}'139 data = {140 'email': user.email,141 'password': password,142 'returnSecureToken': True143 }144 response = requests.post(url, json=data)145 result = response.json()146 if 'idToken' in result:147 st.session_state.logged_in = True148 st.session_state.current_user = user.uid149 st.session_state.display_name = user.display_name # Store the display name150 st.success("Logged in successfully!")151 return152 153 raise Exception("User not found with provided credentials.") # if not found, raise exception.154 except Exception as e:155 st.error(f"Login failed: {e}") # if any error, display this message.156 else:157 raise Exception("Error with sign-in endpoint") # If sign-in endpoint doesn't return error or id token, then throw this error.158 159 except Exception as e:160 st.error(f"Login failed: {e}")161 162# Callback for logout163def logout_callback():164 st.session_state.logged_in = False165 st.session_state.current_user = None166 st.session_state.display_name = None167 st.session_state.selected_image = None168 st.info("Logged out successfully!")169 170# Function to get image from url171def get_image_from_url(url):172 """173 Fetches and returns an image from a given URL, converting to PNG if needed.174 """175 try:176 response = requests.get(url, stream=True)177 response.raise_for_status()178 image = Image.open(BytesIO(response.content))179 return image, url # Return the image and the URL180 except requests.exceptions.RequestException as e:181 return f"Error fetching image: {e}", None182 except Exception as e:183 return f"Error processing image: {e}", None184 185# Function to generate image186def generate_image(prompt, aspect_ratio, realism):187 global token_id188 global TOKEN189 global tokens_tried190 global no_of_accounts191 global model_id192 payload = {193 "id": model_id,194 "inputs": [prompt, aspect_ratio, str(realism).lower()],195 }196 headers = {"Authorization": f"Bearer {TOKEN}"}197 198 try:199 response_data = requests.post(API_URL, json=payload, headers=headers).json()200 if "error" in response_data:201 if 'error 429' in response_data['error']:202 if tokens_tried < no_of_accounts:203 token_id = (token_id + 1) % (no_of_accounts)204 tokens_tried += 1205 TOKEN = os.getenv(f"TOKEN{token_id}")206 response_data = generate_image(prompt, aspect_ratio, realism)207 tokens_tried = 0208 return response_data209 return "No credits available", None210 return response_data, None211 elif "output" in response_data:212 url = response_data['output']213 image, url = get_image_from_url(url)214 return image, url # Return the image and the URL215 else:216 return "Error: Unexpected response from server", None217 except Exception as e:218 return f"Error", None219 220def download_image(image_url):221 if not image_url:222 return None # Return None if image_url is empty223 try:224 response = requests.get(image_url, stream=True)225 response.raise_for_status()226 227 # Get the content type from the headers228 content_type = response.headers.get('content-type')229 extension = mimetypes.guess_extension(content_type)230 231 if not extension:232 extension = ".png" # Default to .png if can't determine the extension233 234 # Create a temporary file with the correct extension235 with tempfile.NamedTemporaryFile(suffix=extension, delete=False) as tmp_file:236 for chunk in response.iter_content(chunk_size=8192):237 tmp_file.write(chunk)238 temp_file_path = tmp_file.name239 return temp_file_path240 except Exception as e:241 return None242 243# Function to store image and related data in Firebase244def store_image_data_in_db(user_id, prompt, aspect_ratio, realism, image_url, thumbnail_url):245 try:246 ref = db.reference(f'users/{user_id}/images')247 new_image_ref = ref.push()248 new_image_ref.set({249 'prompt': prompt,250 'aspect_ratio': aspect_ratio,251 'realism': realism,252 'image_url': image_url,253 'thumbnail_url' : thumbnail_url,254 'timestamp': {'.sv': 'timestamp'}255 })256 st.success("Image data saved successfully!")257 except Exception as e:258 st.error(f"Failed to save image data: {e}")259 260#Function to upload image to cloud storage261def upload_image_to_storage(image, user_id, is_thumbnail = False):262 try:263 bucket = storage.bucket()264 image_id = str(uuid.uuid4())265 if is_thumbnail:266 file_path = f"user_images/{user_id}/thumbnails/{image_id}.png" # path for thumbnail267 else:268 file_path = f"user_images/{user_id}/{image_id}.png" # path for high resolution images269 blob = bucket.blob(file_path)270 271 # Convert PIL Image to BytesIO object272 img_byte_arr = BytesIO()273 image.save(img_byte_arr, format='PNG')274 img_byte_arr = img_byte_arr.getvalue()275 276 blob.upload_from_string(img_byte_arr, content_type='image/png')277 blob.make_public()278 image_url = blob.public_url279 return image_url280 except Exception as e:281 st.error(f"Failed to upload image to cloud storage: {e}")282 return None283 284#Function to load image data from the database285def load_image_data(user_id, start_index, batch_size):286 try:287 ref = db.reference(f'users/{user_id}/images')288 snapshot = ref.order_by_child('timestamp').limit_to_last(start_index + batch_size).get()289 if snapshot:290 image_list = list(snapshot.items())291 image_list.reverse() # Reverse to show latest first292 293 new_images = []294 for key, val in image_list[start_index:]:295 new_images.append(val)296 return new_images297 else:298 return []299 except Exception as e:300 st.error(f"Failed to fetch image data from database: {e}")301 return []302 303# Function to create low resolution thumbnail304def create_thumbnail(image, thumbnail_size = (150,150)):305 try:306 img_byte_arr = BytesIO()307 image.thumbnail(thumbnail_size)308 image.save(img_byte_arr, format='PNG')309 img_byte_arr = img_byte_arr.getvalue()310 thumbnail = Image.open(io.BytesIO(img_byte_arr)) # convert byte to PIL image311 return thumbnail312 except Exception as e:313 st.error(f"Failed to create thumbnail: {e}")314 return None315 316# Registration form317def registration_form():318 with st.form("Registration"):319 st.subheader("Register")320 email = st.text_input("Email", key="reg_email")321 display_name = st.text_input("Display Name", key="reg_display_name")322 password = st.text_input("Password (min 6 characters)", type="password", key="reg_password")323 submit_button = st.form_submit_button("Register", on_click=register_callback)324 325# Login form326def login_form():327 with st.form("Login"):328 st.subheader("Login")329 login_identifier = st.text_input("Email or Username", key="login_identifier")330 password = st.text_input("Password", type="password", key="login_password")331 submit_button = st.form_submit_button("Login", on_click=login_callback)332 333def main_app():334 st.subheader(f"Welcome, {st.session_state.display_name}!")335 st.write("Enter a prompt below to generate an image.")336 337 # Input fields338 prompt = st.text_input("Prompt", key="image_prompt", placeholder="Describe the image you want to generate")339 aspect_ratio = st.radio(340 "Aspect Ratio",341 options=["1:1", "3:4", "4:3", "9:16", "16:9", "9:21", "21:9"],342 index=5343 )344 realism = st.checkbox("Realism", value=False)345 346 if st.button("Generate Image"):347 if prompt:348 with st.spinner("Generating Image..."):349 image_result = generate_image(prompt, aspect_ratio, realism)350 351 if isinstance(image_result, tuple) and len(image_result) == 2:352 image, image_url = image_result353 if isinstance(image, Image.Image):354 # Define the boundary size355 preview_size = 400356 357 # Get original image dimensions358 original_width, original_height = image.size359 360 # Calculate scaling factor to fit within the boundary361 width_ratio = preview_size / original_width362 height_ratio = preview_size / original_height363 364 scaling_factor = min(width_ratio, height_ratio)365 366 # Calculate new dimensions367 new_width = int(original_width * scaling_factor)368 new_height = int(original_height * scaling_factor)369 370 # Resize the image371 resized_image = image.resize((new_width, new_height), Image.LANCZOS)372 373 # Upload the high-resolution image374 cloud_storage_url = upload_image_to_storage(image, st.session_state.current_user, is_thumbnail=False)375 376 if cloud_storage_url:377 # Create thumbnail from the high-resolution image378 thumbnail = create_thumbnail(image)379 380 if thumbnail:381 # Upload thumbnail to cloud storage and store url382 thumbnail_url = upload_image_to_storage(thumbnail, st.session_state.current_user, is_thumbnail=True)383 384 if thumbnail_url:385 # Store image data in database386 store_image_data_in_db(st.session_state.current_user, prompt, aspect_ratio, realism, cloud_storage_url, thumbnail_url)387 st.success("Image stored to database successfully!")388 with st.container(border=True):389 st.image(resized_image, use_column_width=False) # Display the resized image390 st.write(f"**Prompt:** {prompt}")391 st.write(f"**Aspect Ratio:** {aspect_ratio}")392 st.write(f"**Realism:** {realism}")393 download_path = download_image(image_url)394 if download_path:395 st.download_button(label="Download Image", data = open(download_path, "rb"), file_name = f"image.png", key=f"download_high_res_{uuid.uuid4()}")396 else:397 st.error("Failed to upload thumbnail to cloud storage.")398 else:399 st.error("Failed to create thumbnail")400 else:401 st.error("Failed to upload image to cloud storage.")402 else:403 st.error(f"Image generation failed: {image}")404 405 else:406 st.error(f"Image generation failed: {image_result}")407 else:408 st.warning("Please enter a prompt to generate an image.")409 st.header("Your Generated Images")410 # Initialize the current window, if it doesn't exist in session state411 if "current_window_start" not in st.session_state:412 st.session_state.current_window_start = 0413 414 if "window_size" not in st.session_state:415 st.session_state.window_size = 5 # The number of images to display at a time416 417 if "selected_image" not in st.session_state:418 st.session_state.selected_image = None419 420 # Create left and right arrow buttons421 col_left, col_center, col_right = st.columns([1,8,1])422 423 with col_left:424 if st.button("◀️"):425 st.session_state.current_window_start = max(0, st.session_state.current_window_start - st.session_state.window_size)426 427 with col_right:428 if st.button("▶️"):429 st.session_state.current_window_start += st.session_state.window_size430 431 # Dynamically load images for the window432 all_images = load_image_data(st.session_state.current_user, 0, 1000) # load all images433 434 if all_images:435 num_images = len(all_images)436 437 # Calculate the range for images to display438 start_index = st.session_state.current_window_start439 end_index = min(start_index + st.session_state.window_size, num_images)440 441 images_for_window = all_images[start_index:end_index]442 443 # Setup columns for horizontal slider layout444 num_images_to_display = len(images_for_window)445 446 cols = st.columns(num_images_to_display)447 for i, image_data in enumerate(images_for_window):448 with cols[i]:449 if image_data.get('thumbnail_url') and image_data.get('image_url'):450 if st.button("More", key = f"more_{i}"):451 st.session_state.selected_image = image_data452 st.image(image_data['thumbnail_url'], width = 150) #display thumbnail453 454 else:455 st.image(image_data['image_url'], width = 150)456 st.write(f"**Prompt:** {image_data['prompt']}")457 st.write(f"**Aspect Ratio:** {image_data['aspect_ratio']}")458 st.write(f"**Realism:** {image_data['realism']}")459 st.markdown("---")460 else:461 st.write("No image generated yet!")462 463 # Display modal if an image is selected464 if st.session_state.selected_image:465 with st.container(border = True):466 st.image(st.session_state.selected_image['image_url'], use_column_width=True)467 st.write(f"**Prompt:** {st.session_state.selected_image['prompt']}")468 st.write(f"**Aspect Ratio:** {st.session_state.selected_image['aspect_ratio']}")469 st.write(f"**Realism:** {st.session_state.selected_image['realism']}")470 download_path = download_image(st.session_state.selected_image['image_url'])471 if download_path:472 st.download_button(label="Download Image", data = open(download_path, "rb"), file_name = f"image.png", key=f"download_overlay_{uuid.uuid4()}")473 474 if st.button("Close"):475 st.session_state.selected_image = None # close the modal when "close" is clicked476 477 # Logout button478 if st.button("Logout", on_click=logout_callback):479 pass480if st.session_state.logged_in:481 main_app()482else:483 registration_form()484 login_form()