import gradio as gr
import torch
from diffusers import StableDiffusionPipeline
import random
import numpy as np

# 1. Load a CPU-optimized model
# 'segmind/tiny-sd' is much smaller and faster on CPUs than SDXL
model_id = "segmind/tiny-sd"

# Use float32 because CPU doesn't support float16 well
pipe = StableDiffusionPipeline.from_pretrained(
    model_id, 
    torch_dtype=torch.float32
)
pipe = pipe.to("cpu")

# Optimize for CPU speed
pipe.set_progress_bar_config(disable=True)

MAX_SEED = np.iinfo(np.int32).max

def infer(prompt, seed, randomize_seed, width, height):
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    
    generator = torch.Generator("cpu").manual_seed(seed)
    
    # We use very low steps (10-15) because CPU is slow
    image = pipe(
        prompt=prompt,
        generator=generator,
        num_inference_steps=15, 
        guidance_scale=7.0,
        width=width,
        height=height
    ).images[0]
    
    return image, seed

# Simple UI
with gr.Blocks() as demo:
    gr.Markdown("# CodeIgnite CPU Image Engine")
    with gr.Column():
        prompt = gr.Textbox(label="Prompt", placeholder="A simple cat drawing")
        run_button = gr.Button("Generate (CPU Mode)")
        result = gr.Image(label="Result")
        
        with gr.Accordion("Settings", open=False):
            seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
            width = gr.Slider(label="Width", minimum=256, maximum=512, step=32, value=384)
            height = gr.Slider(label="Height", minimum=256, maximum=512, step=32, value=384)

    run_button.click(
        fn=infer,
        inputs=[prompt, seed, randomize_seed, width, height],
        outputs=[result, seed],
        api_name="predict"
    )

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