app.py

import gradio as gr
from diffusers import StableDiffusionXLPipeline, UNet2DConditionModel
import torch
import gc
import time
import os

# --- Конфигурация SDXL-Lightning ---
BASE_MODEL_ID = "stabilityai/stable-diffusion-xl-base-1.0"
LIGHTNING_MODEL_ID = "ByteDance/SDXL-Lightning"
LIGHTNING_CHECKPOINT = "sdxl_lightning_4step_unet.safetensors"
DEVICE = "cpu"
DTYPE = torch.float32

print(f"INFO: Starting with SDXL-Lightning on {DEVICE}")

# --- Загрузка модели ---
pipe = None
try:
    print("INFO: Loading SDXL-Lightning model...")
    load_start = time.time()
    
    # Загружаем базовую модель SDXL
    pipe = StableDiffusionXLPipeline.from_pretrained(
        BASE_MODEL_ID,
        torch_dtype=DTYPE,
        use_safetensors=True,
        variant="fp16",  # Используем fp16 вариант для экономии памяти
        low_cpu_mem_usage=True,
        safety_checker=None
    )
    
    # Загружаем оптимизированный UNet из SDXL-Lightning
    lightning_unet = UNet2DConditionModel.from_pretrained(
        LIGHTNING_MODEL_ID,
        subfolder="unet",
        torch_dtype=DTYPE,
        use_safetensors=True,
        variant=LIGHTNING_CHECKPOINT,
        low_cpu_mem_usage=True
    )
    
    # Заменяем UNet в пайплайне на Lightning версию
    pipe.unet = lightning_unet
    
    # Оптимизации для CPU
    pipe.enable_attention_slicing()
    
    load_end = time.time()
    print(f"INFO: Model loaded successfully in {load_end - load_start:.2f} seconds")
    
except Exception as e:
    print(f"ERROR: Failed to load model: {e}")

# --- Функция генерации промпта ---
def create_prompt_and_negative(user_description, style_choice):
    base_description = f"emoji of {user_description}, icon"
    style_keywords = ""
    negative_prompt = "ugly, tiling, poorly drawn, out of frame, extra limbs, disfigured, deformed, body out of frame, blurry, bad anatomy, blurred, watermark, grainy, signature, cut off, draft, text, words, low quality, lowres, multiple emojis, multiple objects, borders, frame, artifacts, noise, text, letters, signature, username, complex background"

    if style_choice == "Apple (iOS) Style":
        style_keywords = "apple emoji style, ios emoji, 3D icon, glossy finish, smooth gradients, vibrant colors, rounded, clean vector art look, product render, high detail, playful, expressive, centered, iconic, simple background"
        negative_prompt += ", flat, 2d, sketch, hand-drawn, dark, dull, messy, text"
    elif style_choice == "Google (Noto) Style":
        style_keywords = "google emoji style, noto color emoji, flat 2d design, bold black outline, simple geometric shapes, minimalist, cartoonish, expressive, clean lines, vibrant solid colors, centered, iconic, simple background"
        negative_prompt += ", 3d, realistic, shadows, complex textures, messy, text"
    elif style_choice == "Hand-Drawn Style":
        style_keywords = "hand-drawn emoji, doodle style, simple line art, sketch, minimalist drawing, black and white with color accents, MS Paint style, childlike drawing, centered, iconic, simple background"
        negative_prompt += ", 3d, realistic, complex, detailed, professional, polished, text"
    elif style_choice == "Anime Style":
        style_keywords = "anime emoji, manga style, kawaii, cute japanese anime, big eyes, simple shapes, colorful, 2D, flat colors, centered, iconic, simple background"
        negative_prompt += ", realistic, 3d, photorealistic, complex, detailed, western, text"

    final_prompt = f"{base_description}, {style_keywords}, high quality, centered composition"
    return final_prompt, negative_prompt

# --- Функция генерации ---
def generate_emoji(description, style, guidance_scale_val, seed_val, size_choice):
    if pipe is None:
        raise gr.Error("Model not loaded. Please check logs.")
    if not description.strip():
        raise gr.Error("Please enter a description!")

    # Размер
    if size_choice == "Small (faster)":
        height, width = 512, 512  # SDXL лучше работает с размерами кратными 512
    elif size_choice == "Medium":
        height, width = 768, 768
    else:
        height, width = 1024, 1024

    # Сид
    if seed_val is None or seed_val == -1:
        seed = torch.randint(0, 2**32 - 1, (1,)).item()
    else:
        seed = int(seed_val)
    
    generator = torch.Generator(device='cpu').manual_seed(seed)

    # Промпт
    prompt, negative_prompt = create_prompt_and_negative(description, style)
    print(f"INFO: Generating: '{description}' | Style: '{style}' | Size: {width}x{height}")
    print(f"INFO: Guidance: {guidance_scale_val}, Seed: {seed}")
    
    # Генерация
    gc.collect()
    start_time = time.time()
    
    try:
        with torch.no_grad():
            # SDXL-Lightning оптимизирован для 4 шагов
            image = pipe(
                prompt,
                negative_prompt=negative_prompt,
                num_inference_steps=4,  # Lightning оптимизирован для 4 шагов
                guidance_scale=float(guidance_scale_val),
                generator=generator,
                height=height,
                width=width
            ).images[0]
        
        end_time = time.time()
        generation_time = end_time - start_time
        print(f"INFO: Generated in {generation_time:.2f} seconds")
        
        gc.collect()
        return image, seed, f"Generated in {generation_time:.1f} seconds"
    
    except Exception as e:
        print(f"ERROR: Generation failed: {e}")
        gc.collect()
        raise gr.Error(f"Generation failed: {e}")

# --- Интерфейс ---
css = """
#title_custom { 
    text-align: center; 
    margin-bottom: 10px; 
    font-family: 'Poppins', sans-serif; 
    background: linear-gradient(90deg, #FF8C00, #FF2D55);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
}
.container {
    max-width: 900px;
    margin: 0 auto;
}
.gr-button-primary { 
    background: linear-gradient(90deg, #FF8C00, #FF2D55) !important; 
    border: none !important;
}
.gr-button-primary:hover {
    transform: translateY(-2px);
    box-shadow: 0 5px 15px rgba(255, 140, 0, 0.3) !important;
}
.emoji-preview {
    border-radius: 12px;
    overflow: hidden;
    box-shadow: 0 4px 20px rgba(0,0,0,0.1);
}
footer {visibility: hidden}
"""

with gr.Blocks(css=css, theme=gr.themes.Soft(primary_hue=gr.themes.colors.orange)) as demo:
    gr.Markdown("<h1 id='title_custom'>✨ Dreamoji AI Studio 2025 ✨</h1>", elem_id="title_custom")
    gr.Markdown(
        "<div class='container'>"
        "TEST SPACE Turn your ideas into emojis with SDXL-Lightning - ultra-slow generation in just 4 steps! "
        "Describe your emoji, choose a style, and watch the magic happen."
        "</div>"
    )

    with gr.Row():
        with gr.Column(scale=2):
            description_input = gr.Textbox(
                label="1. Describe your Emoji:",
                placeholder="e.g., 'a happy cat', 'a robot with a lightbulb', 'an avocado dancing'"
            )
            style_input = gr.Radio(
                ["Apple (iOS) Style", "Google (Noto) Style", "Hand-Drawn Style", "Anime Style"],
                label="2. Choose Style:",
                value="Apple (iOS) Style"
            )
            
            size_choice = gr.Radio(
                ["Small (faster)", "Medium", "Large (slower)"],
                label="3. Choose Size:",
                value="Small (faster)"
            )
            
            with gr.Accordion("⚙️ Advanced Settings", open=False):
                guidance_scale_slider = gr.Slider(
                    minimum=1.0, 
                    maximum=10.0, 
                    value=7.0, 
                    step=0.5, 
                    label="Guidance Scale",
                    info="How strictly to follow prompt"
                )
                seed_input = gr.Number(
                    label="Seed (-1 for random)", 
                    value=-1
                )

            generate_button = gr.Button("✨ Generate My Dreamoji! ✨", variant="primary")

        with gr.Column(scale=1):
            output_image = gr.Image(
                label="Your Dreamoji:", 
                type="pil", 
                elem_classes="emoji-preview"
            )
            used_seed_output = gr.Textbox(label="Seed:", interactive=False)
            generation_time = gr.Textbox(label="Time:", interactive=False)

    generate_button.click(
        generate_emoji,
        inputs=[description_input, style_input, guidance_scale_slider, seed_input, size_choice],
        outputs=[output_image, used_seed_output, generation_time],
    )

    gr.Markdown("---")
    
    with gr.Accordion("🤔 Tips", open=True):
        gr.Markdown(
            "- **Ultra-Fast Generation:** SDXL-Lightning generates in just 4 steps!\n"
            "- **Keep it simple:** Short descriptions work best\n"
            "- **Try different styles:** Each style has a unique look\n"
            "- **Save the seed:** Use it again to recreate similar results"
        )
    
    gr.HTML("""
        <div style="text-align: center; margin-top: 20px; font-size: 0.9em; color: #777;">
            Created with Gradio & Hugging Face Spaces. Model: SDXL-Lightning by ByteDance.
            <br>
            If you enjoy this, consider supporting by giving the Space a ❤️!
        </div>
    """)

# Запускаем с очередью
demo.queue(max_size=5).launch()