# app.py for Hugging Face Space (Stable Diffusion)
import gradio as gr
from diffusers import StableDiffusionPipeline
import torch
from PIL import Image
import io
import base64

# Global variables for model and device
sd_pipeline = None
device = "cuda" if torch.cuda.is_available() else "cpu"

# Function to load the model (called once at startup)
def load_sd_model():
    global sd_pipeline
    if sd_pipeline is None:
        print(f"--- Loading Stable Diffusion v1.5 model on {device} ---")
        sd_model_id = "runwayml/stable-diffusion-v1-5"
        try:
            # Load in float16 for GPU to save VRAM, or float32 for CPU
            sd_pipeline = StableDiffusionPipeline.from_pretrained(
                sd_model_id, 
                torch_dtype=torch.float16 if device == "cuda" else torch.float32, 
                use_safetensors=True 
            )
            # Move to device explicitly if device_map is not used (for pipeline.to(device))
            sd_pipeline.to(device) 
            
            if device == "cuda":
                sd_pipeline.enable_attention_slicing()
            
            print("Stable Diffusion model loaded successfully!")
        except Exception as e:
            print(f"Error loading Stable Diffusion model: {e}")
            sd_pipeline = None
            raise RuntimeError(f"Failed to load Stable Diffusion model: {e}")
    return sd_pipeline

# Function to generate an image
def generate_image(prompt):
    global sd_pipeline
    if sd_pipeline is None:
        return "Error: Stable Diffusion model not loaded."

    if not prompt:
        return "Please provide a text description for the image."
    
    print(f"--- Generating image for prompt: '{prompt}' ---")
    try:
        # Use num_inference_steps=25 for speed, can be increased for quality
        image = sd_pipeline(prompt, num_inference_steps=25).images[0]
        print("Image generated successfully!")
        return image
    except Exception as e:
        print(f"Error generating image: {e}")
        return f"Failed to generate image: {str(e)}"

# Load the model during startup of the Space
load_sd_model()

# Create Gradio interface
iface = gr.Interface(
    fn=generate_image, 
    inputs=gr.Textbox(lines=2, label="وصف الصورة (باللغة الإنجليزية)"), 
    outputs=gr.Image(type="pil", label="الصورة المولدة"),
    title="مولد الصور Pi-1 (Stable Diffusion)",
    description="أدخل وصفًا نصيًا باللغة الإنجليزية لتوليد صورة."
)

iface.launch(share=False)