Update app.py

app.py

@@ -1,8 +1,10 @@
 import os
+import json
 import gradio as gr
 import vertexai
 from vertexai.preview.vision_models import ImageGenerationModel
-from huggingface_hub import login # Added to make the login function work
+from huggingface_hub import login
+from google.oauth2 import service_account # Added for secure in-memory auth
 
 # --- 1. Configuration and Authentication ---
 
@@ -23,11 +25,8 @@ else:
 # Part B: Google Cloud Credentials and Initialization
 creds_json_str = os.environ.get("GOOGLE_APPLICATION_CREDENTIALS_JSON")
 
-# <<< START: SECURITY UPDATE >>>
 # This global variable will hold a generic error message for the UI if initialization fails.
-# We no longer include raw exception details here.
 generic_init_error_message = "FATAL: The image generation service is not configured correctly. Please contact the administrator."
-# <<< END: SECURITY UPDATE >>>
 
 # Check if all necessary secrets are loaded
 if not all([GCP_PROJECT_ID, GCP_LOCATION, creds_json_str]):
@@ -47,27 +46,31 @@ else:
     # This block runs only if all secrets are present.
     print("All required GCP secrets (Project, Location, Credentials) are loaded.")
     
-    # <<< START: SECURITY UPDATE >>>
-    # Note: Writing credentials to a file is a potential security risk if the container's filesystem is compromised.
-    # This is a common pattern required by some Google Cloud libraries. Ensure the execution environment is secure.
-    # <<< END: SECURITY UPDATE >>>
-    with open("gcp_creds.json", "w") as f:
-        f.write(creds_json_str)
-    os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = "gcp_creds.json"
-    
     try:
-        vertexai.init(project=GCP_PROJECT_ID, location=GCP_LOCATION)
+        # <<< START: SECURITY FIX >>>
+        # Securely initialize the Vertex AI client by passing credentials in-memory.
+        # This avoids writing sensitive credential files to the container's filesystem.
+        
+        # 1. Parse the JSON string from the environment variable into a Python dict.
+        creds_info = json.loads(creds_json_str)
+        
+        # 2. Create a credentials object from the dictionary.
+        credentials = service_account.Credentials.from_service_account_info(creds_info)
+        
+        # 3. Initialize Vertex AI with the in-memory credentials object.
+        vertexai.init(project=GCP_PROJECT_ID, location=GCP_LOCATION, credentials=credentials)
+        # <<< END: SECURITY FIX >>>
+
         generation_model = ImageGenerationModel.from_pretrained("imagen-4.0-generate-preview-06-06")
         print("Vertex AI and Imagen Model initialized successfully.")
         INITIALIZATION_SUCCESS = True
+        
     except Exception as e:
-        # <<< START: SECURITY UPDATE >>>
         # Log the detailed, sensitive error for the administrator ONLY.
         print(f"ERROR: Failed to initialize Vertex AI or the model. Exception: {e}")
         # Set the generic error message for the UI. DO NOT expose the raw exception 'e'.
         error_message = generic_init_error_message
         INITIALIZATION_SUCCESS = False
-        # <<< END: SECURITY UPDATE >>>
 
 # --- 2. The Core Image Generation Function ---
 def generate_image(prompt: str, negative_prompt: str, seed: int):