import gradio as gr
import torch
from diffusers import AutoPipelineForText2Image

# 2. Load the model
# We remove the .to("cuda") from the end of this line.
print("Loading model...")
pipe = AutoPipelineForText2Image.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", 
    torch_dtype=torch.float16, 
    variant="fp16", 
    use_safetensors=True
)
print("Model partially loaded.")

# ### CHANGE HERE ###
# This is the magic line that enables memory-efficient offloading.
pipe.enable_model_cpu_offload()
print("Model loaded with CPU offloading.")


# 3. Define the image generation function (This part remains the same)
def generate_image(prompt, negative_prompt, steps, guidance):
    print(f"Generating image for prompt: {prompt}")
    # The 'pipe' object does all the work. We pass it the prompt and other parameters.
    # num_inference_steps controls how many steps the model takes to generate the image.
    # guidance_scale controls how much the model follows the prompt.
    image = pipe(
        prompt=prompt, 
        negative_prompt=negative_prompt,
        num_inference_steps=steps,
        guidance_scale=guidance
    ).images[0]
    
    # The output is a PIL Image object, which Gradio can display directly.
    return image

# 4. Create the Gradio interface (This part remains the same)
with gr.Blocks(theme=gr.themes.Default(primary_hue="blue", secondary_hue="sky")) as demo:
    gr.Markdown("# 🖼️ Stable Diffusion XL Text-to-Image")
    gr.Markdown("Enter a text prompt and see the magic of AI-powered image generation!")
    
    with gr.Row():
        with gr.Column(scale=4):
            prompt_input = gr.Textbox(label="Your Prompt", placeholder="An astronaut riding a horse on Mars, cinematic, 4k")
            negative_prompt_input = gr.Textbox(label="Negative Prompt", placeholder="low quality, blurry, watermark, text")
            submit_btn = gr.Button("Generate Image", variant="primary")
        
        with gr.Column(scale=1):
            steps_slider = gr.Slider(minimum=10, maximum=100, value=25, step=1, label="Inference Steps")
            guidance_slider = gr.Slider(minimum=0, maximum=20, value=7.5, step=0.1, label="Guidance Scale")
    
    output_image = gr.Image(label="Generated Image", width=768, height=768)

    gr.Examples(
        examples=[
            ["A majestic lion wearing a crown, photorealistic", "cartoon, drawing", 30, 8],
            ["A beautiful cityscape at sunset, painted by Van Gogh", "blurry, modern", 25, 7.5],
            ["A cute robot serving coffee in a futuristic cafe, 3D render", "text, humans", 35, 9],
        ],
        inputs=[prompt_input, negative_prompt_input, steps_slider, guidance_slider]
    )

    # 5. Connect the button to the function (This part remains the same)
    submit_btn.click(
        fn=generate_image,
        inputs=[prompt_input, negative_prompt_input, steps_slider, guidance_slider],
        outputs=output_image
    )

# 6. Launch the application
demo.launch()