Update app.py

app.py

@@ -1,290 +1,270 @@
-# Import necessary libraries
+import gradio as gr
 import torch
-from diffusers import StableDiffusionPipeline, EulerAncestralDiscreteScheduler
-from fastapi import FastAPI, HTTPException
-from fastapi.middleware.cors import CORSMiddleware
-from pydantic import BaseModel
-import io
-import base64
+from diffusers import StableDiffusionPipeline
 from PIL import Image
-import time
-from datetime import datetime
+import io
+import requests
 import os
-from fastapi import Request
-from fastapi.responses import HTMLResponse
-
-
-# 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")
-
-# Add CORS middleware to allow requests from n8n
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
-
-# Force CPU usage
-device = "cpu"
-print(f"Using device: {device}")
-
-# Load model
-model_id = "stabilityai/stable-diffusion-2-1"
-print("Loading pipeline... This may take a few minutes.")
+from datetime import datetime
+import re
+import tempfile
 
+# Try to import your existing OCI connector for direct access
 try:
-    pipe = StableDiffusionPipeline.from_pretrained(
-        model_id,
-        torch_dtype=torch.float32,
-        use_safetensors=True,
-        safety_checker=None,
-        requires_safety_checker=False
-    )
-    pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
-    pipe = pipe.to(device)
-    print("Model loaded successfully on CPU!")
-except Exception as e:
-    print(f"Error loading model: {e}")
-    # Fallback
-    model_id = "dreamlike-art/dreamlike-diffusion-1.0"
-    pipe = StableDiffusionPipeline.from_pretrained(
-        model_id,
-        torch_dtype=torch.float32,
-        use_safetensors=True,
-        safety_checker=None,
-        requires_safety_checker=False
-    )
-    pipe = pipe.to(device)
-    print(f"Fell back to {model_id}")
+    # This will work if we're in the same app context
+    from app import oci_connector
+    DIRECT_OCI_ACCESS = True
+    print("✅ Direct OCI access available - using existing OCI connector")
+except ImportError:
+    DIRECT_OCI_ACCESS = False
+    print("⚠️  Direct OCI access not available - using API endpoint")
 
-# Google Drive Setup
-def setup_google_drive():
-    """Initialize Google Drive service"""
+# Initialize the Stable Diffusion model
+def load_model():
+    """Load and return the Stable Diffusion model"""
+    print("🔄 Loading Stable Diffusion model...")
     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
-        
+        pipe = StableDiffusionPipeline.from_pretrained(
+            "runwayml/stable-diffusion-v1-5", 
+            torch_dtype=torch.float32,
+            safety_checker=None,  # Disable for better performance
+            requires_safety_checker=False
+        ).to("cpu")
+        print("✅ Model loaded successfully!")
+        return pipe
     except Exception as e:
-        print(f"Google Drive setup failed: {e}")
-        return None
+        print(f"❌ Model loading failed: {e}")
+        raise e
 
-# Google Drive Setup with DEBUGGING
-def setup_google_drive():
-    """Initialize Google Drive service with detailed debugging"""
+# Load the model once at startup
+pipe = load_model()
+
+def save_to_oci_direct(image, prompt):
+    """Save image using direct OCI connector access"""
     try:
-        print("Setting up Google Drive...")
-        
-        # Get service account credentials from environment variable
-        credentials_json = os.getenv('GOOGLE_SERVICE_ACCOUNT_JSON')
-        if not credentials_json:
-            print("❌ ERROR: GOOGLE_SERVICE_ACCOUNT_JSON environment variable not found")
-            return None
-            
-        print("✅ Found Google service account JSON")
+        # Create temporary file
+        with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as tmp:
+            image.save(tmp, format='PNG')
+            temp_path = tmp.name
         
-        # Get Shared Drive ID
-        SHARED_DRIVE_ID = os.getenv('SHARED_DRIVE_ID')
-        if not SHARED_DRIVE_ID:
-            print("❌ ERROR: SHARED_DRIVE_ID environment variable not set")
-            return None
-            
-        print(f"✅ Shared Drive ID: {SHARED_DRIVE_ID}")
+        # Create organized filename
+        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"story_{timestamp}_{safe_prompt}.png"
         
-        # Parse the JSON credentials
-        try:
-            service_account_info = json.loads(credentials_json)
-            client_email = service_account_info.get('client_email', 'Unknown')
-            print(f"✅ Service account email: {client_email}")
-        except json.JSONDecodeError as e:
-            print(f"❌ ERROR: Invalid JSON in service account credentials: {e}")
-            return None
+        # Use project-based directory structure
+        object_name = f"storybook-generator/childrens-books/{filename}"
         
-        credentials = service_account.Credentials.from_service_account_info(
-            service_account_info,
-            scopes=['https://www.googleapis.com/auth/drive.file']
-        )
+        # Upload using existing OCI connector
+        success, message = oci_connector.upload_file(temp_path, object_name, None)
         
-        # Build the Drive service
-        drive_service = build('drive', 'v3', credentials=credentials)
-        print("✅ Google Drive service initialized successfully")
-        return drive_service
+        # Clean up temporary file
+        os.unlink(temp_path)
         
+        if success:
+            return f"✅ {message}"
+        else:
+            return f"❌ {message}"
+            
     except Exception as e:
-        print(f"❌ Google Drive setup failed: {str(e)}")
-        import traceback
-        traceback.print_exc()
-        return None
-
-# Initialize Google Drive service
-drive_service = setup_google_drive()
-SHARED_DRIVE_ID = os.getenv('SHARED_DRIVE_ID')
+        return f"❌ Direct upload failed: {str(e)}"
 
-def save_to_google_drive(image, prompt):
-    """Save image to Google Drive Shared Drive with detailed debugging"""
-    if not drive_service:
-        print("❌ Google Drive service not available, skipping save")
-        return None
-        
-    if not SHARED_DRIVE_ID:
-        print("❌ Shared Drive ID not configured, skipping save")
-        return None
-        
+def save_to_oci_via_api(image, prompt):
+    """Save image using the OCI API endpoint"""
     try:
-        print(f"🔄 Attempting to save image to Shared Drive: {SHARED_DRIVE_ID}")
-        
-        # 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"
-        
-        print(f"📁 Filename: {filename}")
-        
         # 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': [SHARED_DRIVE_ID]  # Save to Shared Drive
-        }
+        # Create filename
+        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"story_{timestamp}_{safe_prompt}.png"
+        
+        # Your OCI API endpoint URL
+        # Use relative URL since we're in the same space
+        api_url = "/api/upload"
         
-        print(f"📋 File metadata: {file_metadata}")
+        # For Hugging Face deployment, we need to handle different URL formats
+        try:
+            # Try to get the space URL from environment
+            space_name = os.environ.get('SPACE_NAME', 'yukee1992-oci-video-storage')
+            api_url = f"https://{space_name}.hf.space/api/upload"
+        except:
+            # Fallback to relative URL
+            pass
         
-        # Upload to Google Drive with supportsAllDrives=True
-        media = MediaIoBaseUpload(img_bytes, mimetype='image/png', resumable=True)
+        # Prepare form data for API request
+        files = {
+            'file': (filename, img_bytes.getvalue(), 'image/png')
+        }
         
-        print("⬆️  Starting upload to Google Drive...")
-        file = drive_service.files().create(
-            body=file_metadata,
-            media_body=media,
-            supportsAllDrives=True,  # CRITICAL FOR SHARED DRIVES
-            fields='id, webViewLink'
-        ).execute()
+        data = {
+            'project_id': 'storybook-generator',
+            'subfolder': 'childrens-books'
+        }
         
-        drive_link = file.get('webViewLink')
-        print(f"✅ Image saved to Google Drive: {drive_link}")
-        return drive_link
+        # Make the API request
+        response = requests.post(api_url, files=files, data=data)
         
+        if response.status_code == 200:
+            result = response.json()
+            if result['status'] == 'success':
+                return f"✅ {result['message']}"
+            else:
+                return f"❌ API Error: {result.get('message', 'Unknown error')}"
+        else:
+            return f"❌ HTTP Error: {response.status_code} - {response.text}"
+            
     except Exception as e:
-        print(f"❌ Failed to save to Google Drive: {str(e)}")
-        import traceback
-        traceback.print_exc()
-        return None
+        return f"❌ API upload failed: {str(e)}"
 
-# Add this simple OAuth callback handler
-@app.get("/oauth2callback", response_class=HTMLResponse)
-async def oauth2_callback(request: Request):
-    """
-    Simple endpoint to handle OAuth2 redirect and display the authorization code
-    """
-    code = request.query_params.get("code")
-    if code:
-        # Display the code in a simple HTML page
-        html_content = f"""
-        <html>
-        <head><title>Authentication Successful</title></head>
-        <body>
-            <h2>✅ Authentication Successful!</h2>
-            <p>Your authorization code has been received.</p>
-            <p>Please copy this code and paste it back into Termux:</p>
-            <div style="background: #f0f0f0; padding: 10px; border-radius: 5px; word-break: break-all;">
-                <strong>{code}</strong>
-            </div>
-            <p><br>Then press Enter to complete the process.</p>
-        </body>
-        </html>
-        """
-        return HTMLResponse(content=html_content)
+def save_to_oci(image, prompt):
+    """Main function to save image to OCI using best available method"""
+    if DIRECT_OCI_ACCESS:
+        return save_to_oci_direct(image, prompt)
     else:
-        return HTMLResponse(content="<h2>❌ No authorization code received</h2>")
-
-# 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():
-    drive_status = "connected" if drive_service else "disconnected"
-    return {"status": "healthy", "model": model_id, "google_drive": drive_status}
+        return save_to_oci_via_api(image, prompt)
 
-# Main API endpoint
-@app.post("/generate")
-async def generate_image(request: GenerateRequest):
+def generate_storybook_image(prompt):
+    """Generate an image from text prompt and save to OCI"""
     try:
-        # Enhanced prompt
-        enhanced_prompt = f"masterpiece, best quality, 4K, ultra detailed, photorealistic, sharp focus, studio lighting, professional photography, {request.prompt}"
-        negative_prompt = "blurry, low quality, low resolution, watermark, signature, text, ugly, deformed"
+        # Enhance the prompt for better children's book style
+        enhanced_prompt = f"children's book illustration, colorful, whimsical, cute, {prompt}"
         
-        print(f"Generating image for prompt: {enhanced_prompt}")
+        print(f"🎨 Generating image for prompt: {enhanced_prompt}")
         
         # Generate image
         image = pipe(
-            prompt=enhanced_prompt,
-            negative_prompt=negative_prompt,
-            width=request.width,
-            height=request.height,
-            guidance_scale=9.0,
-            num_inference_steps=request.steps,
-            generator=torch.Generator(device=device)
+            enhanced_prompt,
+            num_inference_steps=20,  # Faster generation
+            guidance_scale=7.5
         ).images[0]
         
-        if image.mode != 'RGB':
-            image = image.convert('RGB')
+        print("✅ Image generated successfully!")
         
-        print("Image generated successfully!")
+        # Save to OCI
+        print("💾 Saving image to OCI storage...")
+        save_status = save_to_oci(image, prompt)
+        print(save_status)
         
-        # 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)
+        return image, save_status
         
-        # Convert to base64 for API response
-        buffered = io.BytesIO()
-        image.save(buffered, format="PNG")
-        img_base64 = base64.b64encode(buffered.getvalue()).decode()
+    except Exception as e:
+        error_msg = f"❌ Generation failed: {str(e)}"
+        print(error_msg)
+        return None, error_msg
+
+def batch_generate_storybook(scenes):
+    """Generate multiple images for a storybook"""
+    if not scenes:
+        return [], "❌ Please provide at least one scene"
+    
+    results = []
+    status_messages = []
+    
+    for i, scene in enumerate(scenes):
+        if not scene.strip():
+            continue
+            
+        print(f"📖 Generating scene {i+1}/{len(scenes)}: {scene}")
+        image, status = generate_storybook_image(scene)
         
-        return {
-            "status": "success",
-            "image": f"data:image/png;base64,{img_base64}",
-            "prompt": request.prompt,
-            "google_drive_link": drive_link,
-            "saved_to_drive": drive_link is not None
-        }
+        if image:
+            results.append((f"Scene {i+1}: {scene}", image))
+        status_messages.append(f"Scene {i+1}: {status}")
+    
+    return results, "\n".join(status_messages)
+
+# Create the Gradio interface
+with gr.Blocks(title="Children's Book Illustrator", theme="soft") as demo:
+    gr.Markdown("# 📚 Children's Book Illustrator")
+    gr.Markdown("Generate beautiful storybook images and automatically save them to your OCI storage")
+    
+    with gr.Tab("Single Image Generation"):
+        with gr.Row():
+            with gr.Column():
+                prompt_input = gr.Textbox(
+                    label="Scene Description",
+                    placeholder="Describe a scene for your storybook...\nExample: A dragon reading a book under a magical tree",
+                    lines=3
+                )
+                generate_btn = gr.Button("🎨 Generate Image", variant="primary")
+            
+            with gr.Column():
+                image_output = gr.Image(label="Generated Image", height=400)
+                status_output = gr.Textbox(label="Status", interactive=False)
+    
+    with gr.Tab("Batch Storybook Generation"):
+        with gr.Row():
+            with gr.Column():
+                scenes_input = gr.Textbox(
+                    label="Story Scenes (One per line)",
+                    placeholder="Enter each scene on a separate line...\nExample:\nA brave knight approaches the castle\nThe dragon guards a treasure chest\nChildren celebrating with the villagers",
+                    lines=6
+                )
+                batch_generate_btn = gr.Button("📖 Generate Storybook", variant="primary")
+            
+            with gr.Column():
+                batch_status = gr.Textbox(label="Generation Status", interactive=False, lines=10)
         
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Generation failed: {str(e)}")
+        # Gallery for batch results
+        gallery_output = gr.Gallery(
+            label="Generated Storybook Scenes",
+            columns=2,
+            height=600
+        )
+    
+    with gr.Tab("About & Help"):
+        gr.Markdown("""
+        ## 🎯 How to Use
+        
+        1. **Single Image**: Enter a scene description and click "Generate Image"
+        2. **Storybook**: Enter multiple scenes (one per line) for a complete story
+        3. **Auto-Save**: All images are automatically saved to your OCI storage
+        
+        ## 📁 Storage Location
+        
+        Images are saved to: `storybook-generator/childrens-books/`
+        
+        ## 💡 Prompt Tips
+        
+        - Be descriptive: "A dragon reading a book under a magical tree"
+        - Add style: "watercolor style, cute, whimsical"
+        - Specify characters: "little mouse exploring a giant forest"
+        
+        ## 🔧 Technical Details
+        
+        - Uses Stable Diffusion v1.5
+        - Saves to OCI Object Storage
+        - Automatic image organization
+        """)
+    
+    # Connect buttons to functions
+    generate_btn.click(
+        fn=generate_storybook_image,
+        inputs=prompt_input,
+        outputs=[image_output, status_output]
+    )
+    
+    batch_generate_btn.click(
+        fn=batch_generate_storybook,
+        inputs=scenes_input,
+        outputs=[gallery_output, batch_status]
+    )
+
+# For Hugging Face Spaces deployment
+def get_app():
+    """Return the Gradio app for Hugging Face"""
+    return demo
 
-# Run the app
+# For local testing
 if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    print("🚀 Starting Children's Book Illustrator...")
+    print(f"📦 OCI Access: {'Direct' if DIRECT_OCI_ACCESS else 'API'}")
+    demo.launch(server_name="0.0.0.0", server_port=7860)
+else:
+    # For Hugging Face deployment
+    print("📦 Hugging Face Space detected")
+    print(f"🔧 OCI Access: {'Direct' if DIRECT_OCI_ACCESS else 'API'}")