Update app.py

app.py

@@ -8,6 +8,14 @@ import io
 import base64
 from PIL import Image
 import time
+from datetime import datetime
+import os
+
+# Google Drive imports
+from google.oauth2 import service_account
+from googleapiclient.discovery import build
+from googleapiclient.http import MediaIoBaseUpload
+import json
 
 # Initialize FastAPI
 app = FastAPI(title="Children's Book Illustrator API")
@@ -53,17 +61,87 @@ except Exception as e:
     pipe = pipe.to(device)
     print(f"Fell back to {model_id}")
 
+# Google Drive Setup
+def setup_google_drive():
+    """Initialize Google Drive service"""
+    try:
+        # Get service account credentials from environment variable
+        credentials_json = os.getenv('GOOGLE_SERVICE_ACCOUNT_JSON')
+        if not credentials_json:
+            print("Google Drive: No service account credentials found")
+            return None
+            
+        # Parse the JSON credentials
+        service_account_info = json.loads(credentials_json)
+        credentials = service_account.Credentials.from_service_account_info(
+            service_account_info,
+            scopes=['https://www.googleapis.com/auth/drive.file']
+        )
+        
+        # Build the Drive service
+        drive_service = build('drive', 'v3', credentials=credentials)
+        print("Google Drive service initialized successfully")
+        return drive_service
+        
+    except Exception as e:
+        print(f"Google Drive setup failed: {e}")
+        return None
+
+# Initialize Google Drive service
+drive_service = setup_google_drive()
+
+def save_to_google_drive(image, prompt):
+    """Save image to Google Drive"""
+    if not drive_service:
+        print("Google Drive not configured, skipping save")
+        return None
+        
+    try:
+        # Create a filename with timestamp
+        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+        safe_prompt = "".join(c for c in prompt[:30] if c.isalnum() or c in (' ', '-', '_')).rstrip()
+        filename = f"storybook_{timestamp}_{safe_prompt}.png"
+        
+        # Convert image to bytes
+        img_bytes = io.BytesIO()
+        image.save(img_bytes, format='PNG')
+        img_bytes.seek(0)
+        
+        # Create file metadata
+        file_metadata = {
+            'name': filename,
+            'mimeType': 'image/png',
+            'parents': ['root']  # Save to root folder, or specify folder ID
+        }
+        
+        # Upload to Google Drive
+        media = MediaIoBaseUpload(img_bytes, mimetype='image/png', resumable=True)
+        file = drive_service.files().create(
+            body=file_metadata,
+            media_body=media,
+            fields='id, webViewLink'
+        ).execute()
+        
+        print(f"Image saved to Google Drive: {file.get('webViewLink')}")
+        return file.get('webViewLink')
+        
+    except Exception as e:
+        print(f"Failed to save to Google Drive: {e}")
+        return None
+
 # Request model
 class GenerateRequest(BaseModel):
     prompt: str
     width: int = 512
     height: int = 512
     steps: int = 25
+    save_to_drive: bool = True  # New option to control saving
 
 # Health check endpoint
 @app.get("/")
 async def health_check():
-    return {"status": "healthy", "model": model_id}
+    drive_status = "connected" if drive_service else "disconnected"
+    return {"status": "healthy", "model": model_id, "google_drive": drive_status}
 
 # Main API endpoint
 @app.post("/generate")
@@ -91,6 +169,11 @@ async def generate_image(request: GenerateRequest):
         
         print("Image generated successfully!")
         
+        # Save to Google Drive if enabled
+        drive_link = None
+        if request.save_to_drive and drive_service:
+            drive_link = save_to_google_drive(image, request.prompt)
+        
         # Convert to base64 for API response
         buffered = io.BytesIO()
         image.save(buffered, format="PNG")
@@ -99,7 +182,9 @@ async def generate_image(request: GenerateRequest):
         return {
             "status": "success",
             "image": f"data:image/png;base64,{img_base64}",
-            "prompt": request.prompt
+            "prompt": request.prompt,
+            "google_drive_link": drive_link,
+            "saved_to_drive": drive_link is not None
         }
         
     except Exception as e: