import gradio as gr
import torch
from diffusers import StableDiffusionPipeline

# 1. Настройка модели
model_id = "Kolyadual/MicroMacro-GenImage-v1-tiny"

# Убираем use_safetensors=True, так как в репозитории их нет
# Добавляем low_cpu_mem_usage для стабильности на бесплатных тарифах
pipe = StableDiffusionPipeline.from_pretrained(
    model_id, 
    torch_dtype=torch.float32,
    use_safetensors=False, 
    low_cpu_mem_usage=True
)
pipe.to("cpu")

# Оптимизация для CPU
pipe.enable_attention_slicing()

def generate(prompt, steps, guidance):
    # Генерация
    image = pipe(
        prompt=prompt, 
        num_inference_steps=int(steps), 
        guidance_scale=float(guidance)
    ).images[0]
    return image

# 2. Интерфейс Gradio
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🧪 MicroMacro GenImage v1 Tiny")
    gr.Markdown("### Optimized for CPU. Please use **English** prompts.")
    
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(label="Prompt (EN)", placeholder="Alchemy crystal...")
            steps = gr.Slider(1, 25, 12, step=1, label="Steps")
            guidance = gr.Slider(1, 15, 7.5, step=0.5, label="Guidance Scale")
            btn = gr.Button("Generate ✨")
        
        with gr.Column():
            output_img = gr.Image(label="Result")

    # Примеры для быстрой проверки пользователями
    gr.Examples(
        examples=[["mystical potion, alchemy style, glowing", 12, 7.5]],
        inputs=[prompt, steps, guidance]
    )

    btn.click(fn=generate, inputs=[prompt, steps, guidance], outputs=output_img)

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