import os
os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
import uuid
import spaces
import torch
from diffusers import ErnieImagePipeline
from gradio import Server
from gradio.data_classes import FileData

# Optimize for performance if on GPU
torch.set_float32_matmul_precision("high")

# Initialize Pipeline
print("Loading model Baidu/ERNIE-Image-Turbo... this may take a few minutes!", flush=True)
try:
    pipe = ErnieImagePipeline.from_pretrained(
        "Baidu/ERNIE-Image-Turbo",
        torch_dtype=torch.bfloat16,
    )
    print("Model loaded successfully. Moving to CUDA...", flush=True)
    pipe = pipe.to("cuda")
    print("Model is on CUDA. Initializing Server...", flush=True)
except Exception as e:
    print(f"Error during model loading: {e}", flush=True)
    raise

app = Server()

@app.api()
@spaces.GPU(duration=120)
def generate_image(prompt: str, width: int = 1024, height: int = 1024, guidance_scale: float = 1.0, num_inference_steps: int = 8) -> FileData:
    """Generate an image using ERNIE-Image-Turbo."""
    print(f"Endpoint triggered! Prompt: {prompt}, width: {width}, height: {height}", flush=True)
    
    # Run pipeline
    image = pipe(
        prompt=prompt,
        height=height,
        width=width,
        num_inference_steps=num_inference_steps,
        guidance_scale=guidance_scale
    ).images[0]
    
    # Save to a temporary unique file
    os.makedirs("/tmp/ernie_outputs", exist_ok=True)
    out_path = f"/tmp/ernie_outputs/{uuid.uuid4()}.png"
    image.save(out_path)
    
    return FileData(path=out_path)

from fastapi.responses import HTMLResponse

@app.get("/")
async def homepage():
    """Serve the custom frontend HTML."""
    html_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "index.html")
    with open(html_path, "r", encoding="utf-8") as f:
        return HTMLResponse(content=f.read())

app.launch(show_error=True)