Update app.py

app.py

@@ -1,27 +1,73 @@
 import gradio as gr
 import numpy as np
 import random
-
-# import spaces #[uncomment to use ZeroGPU]
-from diffusers import DiffusionPipeline
+import os
+from diffusers import DiffusionPipeline, DPMSolverMultistepScheduler
 import torch
+from PIL import Image
+import time
 
+# Ustawienia środowiska dla lepszej wydajności na CPU
 device = "cuda" if torch.cuda.is_available() else "cpu"
-model_repo_id = "dhead/wai-nsfw-illustrious-sdxl-v140-sdxl"  # Replace to the model you would like to use
+torch.set_grad_enabled(False)  # Wyłącz gradienty dla inferencji
+
+# Optymalizacje dla CPU
+if device == "cpu":
+    os.environ["OMP_NUM_THREADS"] = str(os.cpu_count())
+    torch.set_num_threads(os.cpu_count())
 
+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:
     torch_dtype = torch.float32
+    pipe = DiffusionPipeline.from_pretrained(
+        model_repo_id,
+        torch_dtype=torch_dtype,
+        use_safetensors=True
+    )
 
-pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
+# Optymalizacje potoku
+pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
 pipe = pipe.to(device)
 
+# Dodatkowe optymalizacje dla CPU
+if device == "cpu":
+    pipe.enable_attention_slicing()
+    pipe.enable_sequential_cpu_offload()  # Dla systemów z małą ilością RAM
+
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
+DEFAULT_IMAGE_SIZE = 768  # Zmniejszony domyślny rozmiar dla CPU
 
+def optimize_for_prompt(prompt, width, height):
+    """Automatyczna optymalizacja parametrów na podstawie promptu"""
+    prompt_lower = prompt.lower()
+    
+    # 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
+    
+    # Dostosuj rozmiar na podstawie dostępnej pamięci
+    total_pixels = width * height
+    if total_pixels > 1024 * 1024:
+        width = min(width, 768)
+        height = min(height, 768)
+        steps = min(steps, 20)
+    
+    return steps, width, height
 
-# @spaces.GPU #[uncomment to use ZeroGPU]
 def infer(
     prompt,
     negative_prompt,
@@ -31,110 +77,241 @@ def infer(
     height,
     guidance_scale,
     num_inference_steps,
+    enable_optimizations=True,
     progress=gr.Progress(track_tqdm=True),
 ):
+    start_time = time.time()
+    
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
 
     generator = torch.Generator().manual_seed(seed)
+    
+    # Automatyczne optymalizacje
+    if enable_optimizations:
+        num_inference_steps, width, height = optimize_for_prompt(prompt, width, height)
+    
+    try:
+        image = pipe(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            guidance_scale=guidance_scale,
+            num_inference_steps=num_inference_steps,
+            width=width,
+            height=height,
+            generator=generator,
+        ).images[0]
+        
+        generation_time = time.time() - start_time
+        
+        return image, seed, f"Generation time: {generation_time:.2f}s | Steps: {num_inference_steps} | Size: {width}x{height}"
+    
+    except Exception as e:
+        return None, seed, f"Error: {str(e)}"
 
-    image = pipe(
-        prompt=prompt,
-        negative_prompt=negative_prompt,
-        guidance_scale=guidance_scale,
-        num_inference_steps=num_inference_steps,
-        width=width,
-        height=height,
-        generator=generator,
-    ).images[0]
-
-    return image, seed
+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}"
 
+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
 
 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",
 ]
 
 css = """
 #col-container {
     margin: 0 auto;
-    max-width: 640px;
+    max-width: 800px;
+}
+.gallery-container {
+    display: grid;
+    grid-template-columns: repeat(auto-fill, minmax(200px, 1fr));
+    gap: 10px;
+    margin-top: 20px;
+}
+.gallery-item {
+    border-radius: 8px;
+    overflow: hidden;
+}
+.performance-info {
+    background: #f0f0f0;
+    padding: 10px;
+    border-radius: 5px;
+    margin: 10px 0;
 }
 """
 
-with gr.Blocks(css=css) as demo:
+with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
     with gr.Column(elem_id="col-container"):
-        gr.Markdown(" # Text-to-Image Gradio Template")
-
+        gr.Markdown("""
+        # 🎨 Advanced Text-to-Image Generator
+        *Optimized for CPU performance with 18GB RAM*
+        """)
+        
         with gr.Row():
-            prompt = gr.Text(
-                label="Prompt",
-                show_label=False,
-                max_lines=1,
-                placeholder="Enter your prompt",
-                container=False,
-            )
-
-            run_button = gr.Button("Run", scale=0, variant="primary")
-
-        result = gr.Image(label="Result", show_label=False)
-
-        with gr.Accordion("Advanced Settings", open=False):
-            negative_prompt = gr.Text(
-                label="Negative prompt",
-                max_lines=1,
-                placeholder="Enter a negative prompt",
-                visible=False,
-            )
-
-            seed = gr.Slider(
-                label="Seed",
-                minimum=0,
-                maximum=MAX_SEED,
-                step=1,
-                value=0,
-            )
-
-            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-
-            with gr.Row():
-                width = gr.Slider(
-                    label="Width",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=1024,  # Replace with defaults that work for your model
+            with gr.Column(scale=4):
+                prompt = gr.Text(
+                    label="Prompt",
+                    show_label=False,
+                    max_lines=2,
+                    placeholder="Enter your detailed prompt here...",
+                    container=False,
                 )
+            with gr.Column(scale=1):
+                run_button = gr.Button("Generate 🚀", variant="primary", size="lg")
 
-                height = gr.Slider(
-                    label="Height",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=1024,  # Replace with defaults that work for your model
+        with gr.Row():
+            with gr.Column():
+                result = gr.Image(label="Generated Image", show_label=True, height=512)
+                with gr.Row():
+                    save_btn = gr.Button("💾 Save Image")
+                    clear_btn = gr.Button("🗑️ Clear")
+                
+                performance_info = gr.Textbox(
+                    label="Generation Info",
+                    interactive=False,
+                    max_lines=2
                 )
+            
+            with gr.Column():
+                with gr.Accordion("🎛️ Advanced Settings", open=False):
+                    with gr.Tab("Basic"):
+                        negative_prompt = gr.Text(
+                            label="Negative Prompt",
+                            max_lines=2,
+                            placeholder="What to exclude from the image...",
+                            value="blurry, low quality, distorted"
+                        )
+                        
+                        with gr.Row():
+                            seed = gr.Number(
+                                label="Seed",
+                                value=0,
+                                precision=0
+                            )
+                            randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
+                        
+                        enable_optimizations = gr.Checkbox(
+                            label="Enable Auto-Optimizations", 
+                            value=True,
+                            info="Automatically adjust settings for better performance"
+                        )
 
-            with gr.Row():
-                guidance_scale = gr.Slider(
-                    label="Guidance scale",
-                    minimum=0.0,
-                    maximum=10.0,
-                    step=0.1,
-                    value=4.0,  # Replace with defaults that work for your model
-                )
+                    with gr.Tab("Dimensions & Quality"):
+                        with gr.Row():
+                            width = gr.Slider(
+                                label="Width",
+                                minimum=256,
+                                maximum=MAX_IMAGE_SIZE,
+                                step=32,
+                                value=DEFAULT_IMAGE_SIZE,
+                            )
+                            height = gr.Slider(
+                                label="Height",
+                                minimum=256,
+                                maximum=MAX_IMAGE_SIZE,
+                                step=32,
+                                value=DEFAULT_IMAGE_SIZE,
+                            )
+                        
+                        with gr.Row():
+                            guidance_scale = gr.Slider(
+                                label="Guidance Scale",
+                                minimum=1.0,
+                                maximum=10.0,
+                                step=0.1,
+                                value=5.0,
+                            )
+                            num_inference_steps = gr.Slider(
+                                label="Inference Steps",
+                                minimum=10,
+                                maximum=40,
+                                step=1,
+                                value=20,
+                            )
 
-                num_inference_steps = gr.Slider(
-                    label="Number of inference steps",
-                    minimum=1,
-                    maximum=50,
-                    step=1,
-                    value=24,  # Replace with defaults that work for your model
-                )
+                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"
+                    )
 
-        gr.Examples(examples=examples, inputs=[prompt])
-    gr.on(
+        # Przykłady
+        gr.Examples(
+            examples=examples, 
+            inputs=[prompt],
+            label="Click any example to load it ↑"
+        )
+        
+        # Sekcja informacyjna
+        with gr.Accordion("ℹ️ Usage Tips", open=False):
+            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
+            
+            **Recommended Settings for 18GB RAM:**
+            - Max resolution: 1024x1024
+            - Steps: 15-25
+            - Guidance scale: 4.0-7.0
+            """)
+
+    # Główne zdarzenia
+    run_event = gr.on(
         triggers=[run_button.click, prompt.submit],
         fn=infer,
         inputs=[
@@ -146,9 +323,51 @@ with gr.Blocks(css=css) as demo:
             height,
             guidance_scale,
             num_inference_steps,
+            enable_optimizations,
+        ],
+        outputs=[result, seed, performance_info]
+    )
+    
+    # Zdarzenia dodatkowe
+    save_btn.click(
+        fn=save_image,
+        inputs=[result, prompt, seed],
+        outputs=[performance_info]
+    )
+    
+    clear_btn.click(
+        fn=lambda: [None, 0, "Cleared"],
+        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=[result, seed],
+        outputs=[batch_gallery]
+    )
+    
+    # Automatyczne czyszczenie przy zmianie promptu
+    prompt.change(
+        fn=lambda: [None, 0, "Enter new prompt and click Generate"],
+        outputs=[result, seed, performance_info]
     )
 
 if __name__ == "__main__":
-    demo.launch()
+    # Konfiguracja launch dla lepszej wydajności
+    demo.launch(
+        server_name="0.0.0.0",
+        share=False,
+        show_error=True,
+        max_file_size="100MB"
+    )