import gradio as gr
from diffusers import DiffusionPipeline
import torch
import requests

# Load the diffusion model on CPU
pipe = DiffusionPipeline.from_pretrained("tencent/SRPO")
pipe.to("cpu")

FIREBASE_API_KEY = "YOUR_FIREBASE_API_KEY"

def verify_id_token(id_token):
    """
    Verify Firebase ID token using Firebase REST API
    """
    url = f"https://identitytoolkit.googleapis.com/v1/accounts:lookup?key={FIREBASE_API_KEY}"
    response = requests.post(url, json={"idToken": id_token})
    if response.status_code == 200:
        return response.json()
    else:
        return None

def generate_image_with_auth(prompt, id_token):
    # Verify user
    user_info = verify_id_token(id_token)
    if not user_info:
        return None, "❌ Authentication failed. Please log in with Google."
    
    # Generate image on CPU
    image = pipe(prompt, height=256, width=256, num_inference_steps=25).images[0]
    return image, f"✅ Logged in as {user_info['users'][0]['email']}"

# Gradio interface with two outputs: image + login message
with gr.Blocks() as demo:
    gr.Markdown("## SRPO Diffusion Generator (CPU) with Google Login")
    with gr.Row():
        prompt_input = gr.Textbox(label="Prompt")
        id_token_input = gr.Textbox(label="Firebase ID Token (from Google login)", placeholder="Paste ID token here")
    image_output = gr.Image(label="Generated Image")
    login_status = gr.Textbox(label="Login Status")
    generate_btn = gr.Button("Generate Image")
    
    generate_btn.click(
        fn=generate_image_with_auth,
        inputs=[prompt_input, id_token_input],
        outputs=[image_output, login_status]
    )

if __name__ == "__main__":
    demo.launch()