Update app.py

app.py

@@ -6,6 +6,7 @@ from diffusers import DiffusionPipeline, DPMSolverMultistepScheduler
 import torch
 from PIL import Image
 import time
+import psutil
 
 # Ustawienia środowiska dla lepszej wydajności na CPU
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -15,58 +16,99 @@ torch.set_grad_enabled(False)  # Wyłącz gradienty dla inferencji
 if device == "cpu":
     os.environ["OMP_NUM_THREADS"] = str(os.cpu_count())
     torch.set_num_threads(os.cpu_count())
+    print(f"Using {os.cpu_count()} CPU threads")
 
 model_repo_id = "dhead/wai-nsfw-illustrious-sdxl-v140-sdxl"
 
 # Optymalizacje typu danych
-if torch.cuda.is_available():
-    torch_dtype = torch.float16
-    pipe = DiffusionPipeline.from_pretrained(
-        model_repo_id, 
-        torch_dtype=torch_dtype,
-        use_safetensors=True,
-        variant="fp16"
-    )
-else:
+try:
+    if torch.cuda.is_available():
+        torch_dtype = torch.float16
+        pipe = DiffusionPipeline.from_pretrained(
+            model_repo_id, 
+            torch_dtype=torch_dtype,
+            use_safetensors=True,
+            variant="fp16" if any(f for f in ["fp16", "fp16-safetensors"] if f in model_repo_id) else None
+        )
+    else:
+        torch_dtype = torch.float32
+        pipe = DiffusionPipeline.from_pretrained(
+            model_repo_id,
+            torch_dtype=torch_dtype,
+            use_safetensors=True
+        )
+except Exception as e:
+    print(f"Error loading model: {e}")
+    # Fallback to basic loading
+    pipe = DiffusionPipeline.from_pretrained(model_repo_id)
     torch_dtype = torch.float32
-    pipe = DiffusionPipeline.from_pretrained(
-        model_repo_id,
-        torch_dtype=torch_dtype,
-        use_safetensors=True
-    )
 
 # Optymalizacje potoku
-pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
+try:
+    pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
+except:
+    print("Using default scheduler")
+
 pipe = pipe.to(device)
 
-# Dodatkowe optymalizacje dla CPU
+# Optymalizacje tylko dla CPU
 if device == "cpu":
-    pipe.enable_attention_slicing()
-    pipe.enable_sequential_cpu_offload()  # Dla systemów z małą ilością RAM
+    try:
+        pipe.enable_attention_slicing()
+        print("Attention slicing enabled")
+    except Exception as e:
+        print(f"Could not enable attention slicing: {e}")
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
-DEFAULT_IMAGE_SIZE = 768  # Zmniejszony domyślny rozmiar dla CPU
+DEFAULT_IMAGE_SIZE = 512  # Zmniejszony domyślny rozmiar dla CPU
+
+def get_memory_info():
+    """Pobierz informacje o użyciu pamięci"""
+    memory = psutil.virtual_memory()
+    return {
+        'total': memory.total / (1024**3),
+        'available': memory.available / (1024**3),
+        'used': memory.used / (1024**3),
+        'percent': memory.percent
+    }
 
-def optimize_for_prompt(prompt, width, height):
-    """Automatyczna optymalizacja parametrów na podstawie promptu"""
+def optimize_for_prompt_and_memory(prompt, width, height):
+    """Automatyczna optymalizacja parametrów na podstawie promptu i dostępnej pamięci"""
     prompt_lower = prompt.lower()
+    memory_info = get_memory_info()
+    
+    # Bazowa liczba kroków na podstawie złożoności promptu
+    complex_keywords = ['detailed', 'intricate', 'complex', '8k', 'ultra detailed', 'high detail']
+    simple_keywords = ['simple', 'minimal', 'basic', 'sketch']
+    
+    base_steps = 20
     
-    # Dostosuj liczbę kroków na podstawie złożoności promptu
-    complex_keywords = ['detailed', 'intricate', 'complex', '8k', 'ultra detailed']
     if any(keyword in prompt_lower for keyword in complex_keywords):
-        steps = min(30, 25)  # Więcej kroków dla złożonych promptów
-    else:
-        steps = min(20, 25)  # Mniej kroków dla prostych promptów
+        base_steps = min(25, base_steps + 5)
+    elif any(keyword in prompt_lower for keyword in simple_keywords):
+        base_steps = max(15, base_steps - 5)
     
-    # Dostosuj rozmiar na podstawie dostępnej pamięci
-    total_pixels = width * height
-    if total_pixels > 1024 * 1024:
+    # Dostosuj na podstawie dostępnej pamięci
+    if memory_info['available'] < 4:  # Mniej niż 4GB dostępne
+        base_steps = max(15, base_steps - 5)
+        width = min(width, 512)
+        height = min(height, 512)
+    elif memory_info['available'] < 8:  # Mniej niż 8GB dostępne
+        base_steps = max(18, base_steps - 2)
         width = min(width, 768)
         height = min(height, 768)
-        steps = min(steps, 20)
     
-    return steps, width, height
+    # Ogranicz całkowitą liczbę pikseli
+    total_pixels = width * height
+    if total_pixels > 1024 * 1024:
+        scale_factor = (1024 * 1024) / total_pixels
+        width = int(width * scale_factor ** 0.5)
+        height = int(height * scale_factor ** 0.5)
+        width = (width // 32) * 32  # Zaokrąglij do wielokrotności 32
+        height = (height // 32) * 32
+    
+    return base_steps, width, height
 
 def infer(
     prompt,
@@ -80,7 +122,11 @@ def infer(
     enable_optimizations=True,
     progress=gr.Progress(track_tqdm=True),
 ):
+    if not prompt.strip():
+        return None, 0, "Please enter a prompt"
+    
     start_time = time.time()
+    memory_before = get_memory_info()
     
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
@@ -88,10 +134,19 @@ def infer(
     generator = torch.Generator().manual_seed(seed)
     
     # Automatyczne optymalizacje
+    original_steps = num_inference_steps
+    original_width = width
+    original_height = height
+    
     if enable_optimizations:
-        num_inference_steps, width, height = optimize_for_prompt(prompt, width, height)
+        num_inference_steps, width, height = optimize_for_prompt_and_memory(prompt, width, height)
     
     try:
+        # Sprawdź dostępną pamięć przed generowaniem
+        memory_info = get_memory_info()
+        if memory_info['available'] < 2:  # Mniej niż 2GB dostępne
+            return None, seed, "Error: Not enough memory available. Please try with lower resolution or fewer steps."
+        
         image = pipe(
             prompt=prompt,
             negative_prompt=negative_prompt,
@@ -103,55 +158,66 @@ def infer(
         ).images[0]
         
         generation_time = time.time() - start_time
+        memory_after = get_memory_info()
+        
+        info_text = f"✅ Generation time: {generation_time:.1f}s | "
+        info_text += f"Steps: {num_inference_steps} | "
+        info_text += f"Size: {width}x{height} | "
+        info_text += f"Memory: {memory_after['used']:.1f}GB used"
         
-        return image, seed, f"Generation time: {generation_time:.2f}s | Steps: {num_inference_steps} | Size: {width}x{height}"
+        if enable_optimizations and (original_steps != num_inference_steps or original_width != width or original_height != height):
+            info_text += f" | ⚡ Auto-optimized"
+        
+        return image, seed, info_text
     
+    except torch.cuda.OutOfMemoryError:
+        return None, seed, "❌ CUDA Out of Memory Error. Please reduce image size or steps."
+    except RuntimeError as e:
+        if "out of memory" in str(e).lower():
+            return None, seed, "❌ System Out of Memory Error. Please reduce image size or steps."
+        else:
+            return None, seed, f"❌ Runtime Error: {str(e)}"
     except Exception as e:
-        return None, seed, f"Error: {str(e)}"
+        return None, seed, f"❌ Error: {str(e)}"
 
 def save_image(image, prompt, seed):
     """Zapisz wygenerowany obraz"""
     if image is None:
         return "No image to save"
     
-    timestamp = int(time.time())
-    filename = f"generated_{timestamp}_{seed}.png"
-    
-    # Tworzenie folderu jeśli nie istnieje
-    os.makedirs("generated_images", exist_ok=True)
-    filepath = os.path.join("generated_images", filename)
-    
-    image.save(filepath)
-    
-    # Zapisz metadane
-    metadata_file = f"generated_images/metadata_{timestamp}.txt"
-    with open(metadata_file, "w") as f:
-        f.write(f"Prompt: {prompt}\n")
-        f.write(f"Seed: {seed}\n")
-        f.write(f"Timestamp: {timestamp}\n")
-    
-    return f"Image saved as {filename}"
+    try:
+        timestamp = int(time.time())
+        filename = f"generated_{timestamp}_{seed}.png"
+        
+        # Tworzenie folderu jeśli nie istnieje
+        os.makedirs("generated_images", exist_ok=True)
+        filepath = os.path.join("generated_images", filename)
+        
+        image.save(filepath)
+        
+        # Zapisz metadane
+        metadata_file = f"generated_images/metadata_{timestamp}.txt"
+        with open(metadata_file, "w") as f:
+            f.write(f"Prompt: {prompt}\n")
+            f.write(f"Seed: {seed}\n")
+            f.write(f"Timestamp: {timestamp}\n")
+            f.write(f"Model: {model_repo_id}\n")
+        
+        return f"✅ Image saved as {filename}"
+    except Exception as e:
+        return f"❌ Error saving image: {str(e)}"
 
-def batch_generate(prompts, num_images, **kwargs):
-    """Funkcja do generowania wielu obrazów"""
-    results = []
-    for prompt in prompts.split('\n'):
-        prompt = prompt.strip()
-        if prompt:
-            for i in range(num_images):
-                kwargs['prompt'] = prompt
-                kwargs['randomize_seed'] = True
-                image, seed, info = infer(**kwargs)
-                if image:
-                    results.append((image, prompt, seed, info))
-    return results
+def clear_all():
+    """Wyczyść wszystkie wyniki"""
+    return None, 0, "Ready for new generation"
 
+# Przykłady
 examples = [
-    "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
-    "An astronaut riding a green horse",
-    "A delicious ceviche cheesecake slice",
-    "Cyberpunk cityscape at night, neon lights, rain",
-    "Majestic dragon flying over medieval castle",
+    "A beautiful sunset over mountains, digital art",
+    "A cute cat wearing a wizard hat, fantasy art",
+    "Futuristic city with flying cars, cyberpunk style",
+    "Peaceful forest with glowing mushrooms, magical",
+    "A bowl of fruit on a table, still life painting",
 ]
 
 css = """
@@ -165,15 +231,19 @@ css = """
     gap: 10px;
     margin-top: 20px;
 }
-.gallery-item {
-    border-radius: 8px;
-    overflow: hidden;
-}
 .performance-info {
     background: #f0f0f0;
     padding: 10px;
     border-radius: 5px;
     margin: 10px 0;
+    font-family: monospace;
+}
+.memory-warning {
+    background: #fff3cd;
+    border: 1px solid #ffeaa7;
+    padding: 10px;
+    border-radius: 5px;
+    margin: 10px 0;
 }
 """
 
@@ -181,7 +251,15 @@ with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
     with gr.Column(elem_id="col-container"):
         gr.Markdown("""
         # 🎨 Advanced Text-to-Image Generator
-        *Optimized for CPU performance with 18GB RAM*
+        *Optimized for CPU performance - 18GB RAM*
+        """)
+        
+        # Wyświetl informacje o systemie
+        memory_info = get_memory_info()
+        gr.Markdown(f"""
+        <div class="performance-info">
+        💻 **System Info**: CPU Mode | 🧠 **Memory**: {memory_info['used']:.1f}GB / {memory_info['total']:.1f}GB used ({memory_info['percent']:.1f}%)
+        </div>
         """)
         
         with gr.Row():
@@ -190,7 +268,7 @@ with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
                     label="Prompt",
                     show_label=False,
                     max_lines=2,
-                    placeholder="Enter your detailed prompt here...",
+                    placeholder="Describe the image you want to generate...",
                     container=False,
                 )
             with gr.Column(scale=1):
@@ -198,15 +276,15 @@ with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
 
         with gr.Row():
             with gr.Column():
-                result = gr.Image(label="Generated Image", show_label=True, height=512)
+                result = gr.Image(label="Generated Image", show_label=True, height=400)
                 with gr.Row():
                     save_btn = gr.Button("💾 Save Image")
                     clear_btn = gr.Button("🗑️ Clear")
                 
                 performance_info = gr.Textbox(
-                    label="Generation Info",
+                    label="Generation Information",
                     interactive=False,
-                    max_lines=2
+                    max_lines=3
                 )
             
             with gr.Column():
@@ -216,7 +294,7 @@ with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
                             label="Negative Prompt",
                             max_lines=2,
                             placeholder="What to exclude from the image...",
-                            value="blurry, low quality, distorted"
+                            value="blurry, low quality, distorted, bad anatomy"
                         )
                         
                         with gr.Row():
@@ -230,7 +308,7 @@ with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
                         enable_optimizations = gr.Checkbox(
                             label="Enable Auto-Optimizations", 
                             value=True,
-                            info="Automatically adjust settings for better performance"
+                            info="Automatically adjust settings for better performance and memory usage"
                         )
 
                     with gr.Tab("Dimensions & Quality"):
@@ -256,58 +334,41 @@ with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
                                 minimum=1.0,
                                 maximum=10.0,
                                 step=0.1,
-                                value=5.0,
+                                value=7.0,
                             )
                             num_inference_steps = gr.Slider(
                                 label="Inference Steps",
                                 minimum=10,
-                                maximum=40,
+                                maximum=30,
                                 step=1,
                                 value=20,
                             )
 
-                with gr.Accordion("🔄 Batch Generation", open=False):
-                    batch_prompts = gr.Textbox(
-                        label="Batch Prompts (one per line)",
-                        lines=3,
-                        placeholder="Enter multiple prompts, one per line..."
-                    )
-                    num_images_per_prompt = gr.Slider(
-                        label="Images per prompt",
-                        minimum=1,
-                        maximum=5,
-                        step=1,
-                        value=1
-                    )
-                    batch_button = gr.Button("Generate Batch", variant="secondary")
-                    
-                    batch_gallery = gr.Gallery(
-                        label="Batch Results",
-                        show_label=True,
-                        columns=3,
-                        height="auto"
-                    )
-
         # Przykłady
         gr.Examples(
             examples=examples, 
             inputs=[prompt],
-            label="Click any example to load it ↑"
+            label="Quick Start Examples - Click any example below to load it:"
         )
         
         # Sekcja informacyjna
-        with gr.Accordion("ℹ️ Usage Tips", open=False):
+        with gr.Accordion("ℹ️ Usage Tips & Information", open=True):
             gr.Markdown("""
-            **Performance Tips for CPU:**
-            - Use 512x512 or 768x768 resolutions for faster generation
-            - Keep inference steps between 15-25
-            - Enable auto-optimizations for best results
-            - Use clear, descriptive prompts for better quality
+            **🎯 Performance Tips for CPU (18GB RAM):**
+            - Use **512x512** resolution for fastest generation
+            - **15-25 steps** usually provide good quality
+            - Enable **Auto-Optimizations** for best results
+            - Keep **Guidance Scale** between 5.0-8.0
+            
+            **⚠️ Memory Management:**
+            - Larger images (1024x1024) will use more memory
+            - Complex prompts may require more steps
+            - System automatically optimizes based on available memory
             
-            **Recommended Settings for 18GB RAM:**
-            - Max resolution: 1024x1024
-            - Steps: 15-25
-            - Guidance scale: 4.0-7.0
+            **💡 Prompt Tips:**
+            - Be specific and descriptive
+            - Include style keywords (digital art, painting, photo, etc.)
+            - Use negative prompts to exclude unwanted elements
             """)
 
     # Główne zdarzenia
@@ -336,38 +397,26 @@ with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
     )
     
     clear_btn.click(
-        fn=lambda: [None, 0, "Cleared"],
+        fn=clear_all,
         outputs=[result, seed, performance_info]
     )
     
-    batch_button.click(
-        fn=batch_generate,
-        inputs=[
-            batch_prompts,
-            num_images_per_prompt,
-            negative_prompt,
-            seed,
-            randomize_seed,
-            width,
-            height,
-            guidance_scale,
-            num_inference_steps,
-            enable_optimizations,
-        ],
-        outputs=[batch_gallery]
-    )
-    
     # Automatyczne czyszczenie przy zmianie promptu
     prompt.change(
-        fn=lambda: [None, 0, "Enter new prompt and click Generate"],
+        fn=clear_all,
         outputs=[result, seed, performance_info]
     )
 
 if __name__ == "__main__":
+    print("Starting Text-to-Image Application...")
+    print(f"Device: {device}")
+    print(f"Torch threads: {torch.get_num_threads()}")
+    
     # Konfiguracja launch dla lepszej wydajności
     demo.launch(
         server_name="0.0.0.0",
         share=False,
         show_error=True,
-        max_file_size="100MB"
+        max_file_size="50MB",
+        inbrowser=False
     )