import gradio as gr
import torch
from diffusers import StableDiffusionXLPipeline
import time
import os

# Отключаем autocast для CPU
os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"

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

# Определяем устройство
device = "cpu"

print("🔄 Loading MicroMacro-GenImage-v1-base (SDXL)...")
print("⚡ Optimized for CPU. Generation may take 2-5 minutes.")

# Загружаем модель с правильными настройками для CPU
pipe = StableDiffusionXLPipeline.from_pretrained(
    model_id,
    torch_dtype=torch.float32,  # Используем float32 для CPU
    use_safetensors=True,
    low_cpu_mem_usage=True,
    safety_checker=None,  # Отключаем safety checker для экономии памяти
    requires_safety_checker=False
)

# Перемещаем на CPU
pipe.to(device)

# Оптимизации для CPU
pipe.enable_attention_slicing()  # Уменьшает использование памяти
pipe.enable_vae_slicing()  # Дополнительная экономия памяти для VAE

# Отключаем autocast для CPU
pipe.to = lambda device: None  # Заглушка для предотвращения autocast

print("✅ Model loaded successfully!")

# Функция генерации
def generate(prompt, negative_prompt, steps, guidance, seed):
    # Устанавливаем seed для воспроизводимости
    if seed >= 0:
        generator = torch.Generator(device="cpu").manual_seed(seed)
    else:
        generator = None
    
    start_time = time.time()
    
    try:
        # Прогресс-сообщение
        yield None, f"🔄 Генерация... (шаг 0/{steps})"
        
        # Создаем callback для отслеживания прогресса
        def progress_callback(step, timestep, latents):
            if step % 5 == 0 or step == steps - 1:
                print(f"Step {step + 1}/{steps}")
        
        # Генерация с параметрами SDXL
        result = pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,
            num_inference_steps=int(steps),
            guidance_scale=float(guidance),
            generator=generator,
            width=1024,
            height=1024,
            callback=progress_callback,
            callback_steps=1
        )
        
        image = result.images[0]
        generation_time = time.time() - start_time
        
        yield image, f"✅ Готово! Время генерации: {generation_time:.1f} сек."
        
    except Exception as e:
        yield None, f"❌ Ошибка: {str(e)}"
        print(f"Error details: {e}")

# 2. Интерфейс Gradio
with gr.Blocks(theme=gr.themes.Soft(), title="MicroMacro GenImage v1 Base") as demo:
    gr.HTML("""
    <div style="text-align: center; margin-bottom: 20px;">
        <h1>🎨 MicroMacro GenImage v1 Base</h1>
        <p style="font-size: 18px; color: #555;">
            SDXL-based модель для генерации изображений в стиле MicroMacro<br>
            <small>⚡ Оптимизировано для CPU | ⏱️ Время генерации: 2-5 минут</small>
        </p>
    </div>
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            # Основные параметры
            prompt = gr.Textbox(
                label="📝 Prompt (EN)",
                placeholder="A mystical alchemy laboratory with glowing potions, micro and macro details, intricate textures, 8k, photorealistic",
                lines=3,
                value="A mystical alchemy laboratory with glowing potions, micro and macro details"
            )
            
            negative_prompt = gr.Textbox(
                label="🚫 Negative Prompt",
                placeholder="ugly, blurry, low quality, distorted, deformed",
                lines=2,
                value="ugly, blurry, low quality, distorted, deformed, bad anatomy"
            )
            
            with gr.Row():
                steps = gr.Slider(
                    10, 40, 20, step=1, 
                    label="Steps",
                    info="Рекомендуем 20-25 шагов"
                )
                guidance = gr.Slider(
                    1, 15, 7.5, step=0.5,
                    label="Guidance Scale",
                    info="7.5 - оптимально"
                )
            
            seed = gr.Number(
                label="🎲 Seed (-1 для случайного)",
                value=-1,
                precision=0
            )
            
            generate_btn = gr.Button("✨ Сгенерировать", variant="primary", size="lg")
            
            status = gr.Textbox(label="Статус", interactive=False, value="Готов к работе")
        
        with gr.Column(scale=2):
            output_img = gr.Image(label="Результат", height=512)
    
    # Примеры
    gr.Markdown("### 📋 Примеры промптов")
    examples_data = [
        ["Mystical alchemy circle with glowing runes, micro details of crystals, macro view of potion bottles, dramatic lighting, fantasy art, 8k"],
        ["Microscopic view of organic cellular structures, vibrant colors, macro photography, scientific illustration, detailed textures"],
        ["Steampunk laboratory with brass instruments, macro shot of intricate gears, micro details on vintage equipment, warm lighting"],
        ["Fantasy potion ingredients, macro view of glowing mushrooms and crystals, magical atmosphere, hyperrealistic, detailed"],
    ]
    
    gr.Examples(
        examples=examples_data,
        inputs=[prompt],
        outputs=None,
        label="Нажми на пример для автозаполнения"
    )
    
    # Обработчик генерации с прогрессом
    generate_event = generate_btn.click(
        fn=generate,
        inputs=[prompt, negative_prompt, steps, guidance, seed],
        outputs=[output_img, status]
    )

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