Update app.py

app.py

@@ -1,59 +1,56 @@
 import gradio as gr
-from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
+from diffusers import StableDiffusionXLPipeline, UNet2DConditionModel
 import torch
 import gc
 import time
 import os
 
-# --- Оптимизированная конфигурация ---
-MODEL_ID = "playgroundai/playground-v2.5-1024px-aesthetic"  # Современная модель 2024 года
-DEVICE = "cpu"  # Принудительно CPU для бесплатного тира
-DTYPE = torch.float32  # Для CPU используем float32
+# --- Конфигурация 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: Space starting. Device: {DEVICE}, Model: {MODEL_ID}")
-print(f"INFO: Available RAM: {os.sysconf('SC_PAGE_SIZE') * os.sysconf('SC_PHYS_PAGES') / (1024**3):.2f} GB")
+print(f"INFO: Starting with SDXL-Lightning on {DEVICE}")
 
-# --- Функция для очистки кэша и памяти ---
-def clear_memory():
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
-    gc.collect()
-
-# --- Загрузка модели с оптимизациями ---
+# --- Загрузка модели ---
 pipe = None
 try:
-    print(f"INFO: Loading model {MODEL_ID}...")
-    load_start_time = time.time()
+    print("INFO: Loading SDXL-Lightning model...")
+    load_start = time.time()
     
-    # Создаем оптимизированный планировщик
-    scheduler = DPMSolverMultistepScheduler.from_pretrained(
-        MODEL_ID,
-        subfolder="scheduler",
-        solver_order=2,
-        prediction_type="epsilon",
-        thresholding=False,
-        algorithm_type="sde-dpmsolver++",
-        solver_type="midpoint"
+    # Загружаем базовую модель 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
     )
     
-    # Загружаем модель с оптимизациями
-    pipe = StableDiffusionPipeline.from_pretrained(
-        MODEL_ID,
-        scheduler=scheduler,
+    # Загружаем оптимизированный UNet из SDXL-Lightning
+    lightning_unet = UNet2DConditionModel.from_pretrained(
+        LIGHTNING_MODEL_ID,
+        subfolder="unet",
         torch_dtype=DTYPE,
         use_safetensors=True,
-        low_cpu_mem_usage=True,
-        safety_checker=None  # Отключаем для скорости
+        variant=LIGHTNING_CHECKPOINT,
+        low_cpu_mem_usage=True
     )
     
-    # Применяем оптимизации
-    pipe.enable_attention_slicing(1)
+    # Заменяем UNet в пайплайне на Lightning версию
+    pipe.unet = lightning_unet
+    
+    # Оптимизации для CPU
+    pipe.enable_attention_slicing()
     
-    load_end_time = time.time()
-    print(f"INFO: Model loaded successfully in {load_end_time - load_start_time:.2f} seconds.")
+    load_end = time.time()
+    print(f"INFO: Model loaded successfully in {load_end - load_start:.2f} seconds")
     
 except Exception as e:
-    print(f"ERROR: Critical error loading model: {e}")
+    print(f"ERROR: Failed to load model: {e}")
 
 # --- Функция генерации промпта ---
 def create_prompt_and_negative(user_description, style_choice):
@@ -77,77 +74,64 @@ def create_prompt_and_negative(user_description, style_choice):
     final_prompt = f"{base_description}, {style_keywords}, high quality, centered composition"
     return final_prompt, negative_prompt
 
-# --- Оптимизированная функция генерации ---
-def generate_emoji(description, style, num_steps, guidance_scale_val, seed_val, size_choice):
+# --- Функция генерации ---
+def generate_emoji(description, style, guidance_scale_val, seed_val, size_choice):
     if pipe is None:
-        raise gr.Error("ERROR: Model is not loaded. The Space might be restarting or encountered an issue. Please check logs or try again later.")
+        raise gr.Error("Model not loaded. Please check logs.")
     if not description.strip():
-        raise gr.Error("Please enter a description for your emoji!")
+        raise gr.Error("Please enter a description!")
 
-    # Определяем размер изображения
+    # Размер
     if size_choice == "Small (faster)":
-        height, width = 384, 384
+        height, width = 512, 512  # SDXL лучше работает с размерами кратными 512
     elif size_choice == "Medium":
-        height, width = 512, 512
-    else:  # Large
-        height, width = 640, 640
+        height, width = 768, 768
+    else:
+        height, width = 1024, 1024
 
-    # Генерация сида
-    if seed_val is None or seed_val == -1 or seed_val == 0:
+    # Сид
+    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 for: '{description}' | Style: '{style}' | Size: {width}x{height}")
-    print(f"INFO: Prompt: {prompt}")
-    print(f"INFO: Negative Prompt: {negative_prompt}")
-    print(f"INFO: Steps: {num_steps}, Guidance: {guidance_scale_val}, Seed: {seed}")
-
-    # Ограничение шагов для CPU
-    effective_steps = min(int(num_steps), 20)
-    if int(num_steps) > 20:
-        print(f"WARN: Steps capped at 20 for CPU. Original: {num_steps}")
-
-    generation_start_time = time.time()
+    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:
-        # Очищаем память перед генерацией
-        clear_memory()
-        
-        # Используем torch.no_grad() для экономии памяти
         with torch.no_grad():
+            # SDXL-Lightning оптимизирован для 4 шагов
             image = pipe(
                 prompt,
                 negative_prompt=negative_prompt,
-                num_inference_steps=effective_steps,
+                num_inference_steps=4,  # Lightning оптимизирован для 4 шагов
                 guidance_scale=float(guidance_scale_val),
                 generator=generator,
                 height=height,
                 width=width
             ).images[0]
-
-        generation_end_time = time.time()
-        generation_time = generation_end_time - generation_start_time
-        print(f"INFO: Image generated successfully in {generation_time:.2f} seconds.")
-
-        # Очищаем память после генерации
-        clear_memory()
         
+        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:
-        generation_end_time = time.time()
-        print(f"ERROR: Exception during generation ({generation_end_time - generation_start_time:.2f}s): {e}")
-        clear_memory()
-        
-        if "out of memory" in str(e).lower():
-            raise gr.Error(f"Generation failed: Out of memory. Try a smaller size or fewer steps. (Error: {e})")
-        else:
-            raise gr.Error(f"Generation failed. Please check your input or try again. (Error: {e})")
+        print(f"ERROR: Generation failed: {e}")
+        gc.collect()
+        raise gr.Error(f"Generation failed: {e}")
 
-# --- Современный интерфейс Gradio ---
+# --- Интерфейс ---
 css = """
 #title_custom { 
     text-align: center; 
@@ -169,15 +153,6 @@ css = """
     transform: translateY(-2px);
     box-shadow: 0 5px 15px rgba(255, 140, 0, 0.3) !important;
 }
-.gr-form {
-    border-radius: 12px !important;
-    box-shadow: 0 4px 20px rgba(0,0,0,0.08) !important;
-}
-.gr-success { 
-    background-color: #E6FFFA !important; 
-    color: #1A7F47 !important; 
-    border-color: #A1E0C0 !important; 
-}
 .emoji-preview {
     border-radius: 12px;
     overflow: hidden;
@@ -186,94 +161,81 @@ css = """
 footer {visibility: hidden}
 """
 
-with gr.Blocks(css=css, theme=gr.themes.Soft(primary_hue=gr.themes.colors.orange, secondary_hue=gr.themes.colors.blue)) as demo:
+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'>"
-        "Turn your wildest emoji ideas into reality with Playground v2.5, a modern AI model from late 2024! "
-        "Describe your emoji, pick a style, and watch the magic happen. "
-        "<b>Note:</b> Generation on free CPU hardware takes 1-3 minutes per emoji."
+        "Turn your ideas into emojis with SDXL-Lightning - ultra-fast 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 Dream Emoji:",
-                placeholder="e.g., 'a sleepy sloth drinking coffee', 'a confused robot with a lightbulb', 'a joyful avocado dancing'"
+                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 Emoji Style:",
+                label="2. Choose Style:",
                 value="Apple (iOS) Style"
             )
             
-            with gr.Row():
-                size_choice = gr.Radio(
-                    ["Small (faster)", "Medium", "Large (slower)"],
-                    label="3. Choose Size:",
-                    value="Small (faster)",
-                    info="Smaller = faster generation"
-                )
+            size_choice = gr.Radio(
+                ["Small (faster)", "Medium", "Large (slower)"],
+                label="3. Choose Size:",
+                value="Small (faster)"
+            )
             
             with gr.Accordion("⚙️ Advanced Settings", open=False):
-                with gr.Row():
-                    num_steps_slider = gr.Slider(
-                        minimum=10, 
-                        maximum=25, 
-                        value=15, 
-                        step=1, 
-                        label="Inference Steps",
-                        info="More steps = more detail but slower (max 20 on CPU)"
-                    )
-                    guidance_scale_slider = gr.Slider(
-                        minimum=1.0, 
-                        maximum=12.0, 
-                        value=7.0, 
-                        step=0.5, 
-                        label="Guidance Scale",
-                        info="How strictly to follow prompt"
-                    )
+                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 (use -1 for random, or a number to reproduce results)", 
+                    label="Seed (-1 for random)", 
                     value=-1
                 )
 
-            generate_button = gr.Button("✨ Generate My Dreamoji! ✨", variant="primary", elem_id="generate_button_custom")
+            generate_button = gr.Button("✨ Generate My Dreamoji! ✨", variant="primary")
 
         with gr.Column(scale=1):
             output_image = gr.Image(
-                label="Your Generated Dreamoji:", 
+                label="Your Dreamoji:", 
                 type="pil", 
                 elem_classes="emoji-preview"
             )
-            used_seed_output = gr.Textbox(label="Seed Used:", interactive=False)
-            generation_time = gr.Textbox(label="Generation Time:", interactive=False)
+            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, num_steps_slider, guidance_scale_slider, seed_input, size_choice],
+        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 for Awesome Dreamojis", open=True):
+    with gr.Accordion("🤔 Tips", open=True):
         gr.Markdown(
-            "- **Keep it simple:** 'A happy cat' works better than complex descriptions\n"
-            "- **Try different styles:** Each style has its own unique look\n"
-            "- **Use the seed:** Found a good result? Save the seed and tweak the prompt\n"
-            "- **Size matters:** Smaller sizes generate much faster on CPU\n"
-            "- **Fewer steps:** For quick tests, try 10-15 steps first"
+            "- **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: Playground v2.5 (2024).
+            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=10).launch()
+# Запускаем с очередью
+demo.queue(max_size=5).launch()