Create app.py

app.py

@@ -0,0 +1,166 @@
+import gradio as gr
+from diffusers import StableDiffusionPipeline
+import torch
+from PIL import Image
+import re
+import uuid
+import gc
+
+# Load word list for safety checking (using a simple list instead of loading dataset)
+BLOCKED_WORDS = ["nsfw", "nude", "explicit"]  # Add more as needed
+
+# Initialize the pipeline with CPU optimization
+def initialize_pipeline():
+    pipe = StableDiffusionPipeline.from_pretrained(
+        "stabilityai/stable-diffusion-2-1-base",
+        torch_dtype=torch.float32  # Use float32 for CPU
+    )
+    pipe = pipe.to("cpu")
+    # Enable memory efficient attention
+    pipe.enable_attention_slicing()
+    return pipe
+
+pipe = initialize_pipeline()
+
+def cleanup():
+    """Force garbage collection to free memory"""
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+def infer(prompt, negative, scale):
+    # Safety check
+    prompt = prompt.lower()
+    for word in BLOCKED_WORDS:
+        if word in prompt:
+            raise gr.Error("Unsafe content found. Please try again with different prompts.")
+    
+    try:
+        # Generate only one image at a time to conserve memory
+        images = pipe(
+            prompt=prompt,
+            negative_prompt=negative,
+            guidance_scale=scale,
+            num_inference_steps=30,  # Reduced steps for faster generation
+            num_images_per_prompt=1
+        ).images
+        
+        # Save image
+        output_path = f"{uuid.uuid4()}.jpg"
+        images[0].save(output_path)
+        
+        # Cleanup to free memory
+        cleanup()
+        
+        return [output_path]
+    
+    except Exception as e:
+        cleanup()
+        raise gr.Error(f"Generation failed: {str(e)}")
+
+css = """
+        .gradio-container {
+            max-width: 768px !important;
+            font-family: 'IBM Plex Sans', sans-serif;
+        }
+        .gr-button {
+            color: white;
+            border-color: black;
+            background: black;
+        }
+        input[type='range'] {
+            accent-color: black;
+        }
+        .dark input[type='range'] {
+            accent-color: #dfdfdf;
+        }
+        #gallery {
+            min-height: 22rem;
+            margin-bottom: 15px;
+        }
+        #gallery>div>.h-full {
+            min-height: 20rem;
+        }
+"""
+
+examples = [
+    [
+        'A small cabin on top of a snowy mountain, artstation style',
+        'low quality, ugly',
+        9
+    ],
+    [
+        'A red apple on a wooden table, still life',
+        'low quality',
+        9
+    ],
+]
+
+with gr.Blocks(css=css) as block:
+    gr.HTML(
+        """
+        <div style="text-align: center; margin: 0 auto;">
+            <h1 style="font-weight: 900; margin-bottom: 7px;">
+                Stable Diffusion 2.1 (CPU Version)
+            </h1>
+            <p style="margin-bottom: 10px; font-size: 94%;">
+                Optimized for CPU usage with 16GB RAM
+            </p>
+        </div>
+        """
+    )
+    
+    with gr.Group():
+        with gr.Row():
+            with gr.Column(scale=3):
+                text = gr.Textbox(
+                    label="Enter your prompt",
+                    show_label=False,
+                    max_lines=1,
+                    placeholder="Enter your prompt"
+                )
+                negative = gr.Textbox(
+                    label="Enter your negative prompt",
+                    show_label=False,
+                    max_lines=1,
+                    placeholder="Enter a negative prompt"
+                )
+            with gr.Column(scale=1, min_width=150):
+                btn = gr.Button("Generate image")
+
+        gallery = gr.Gallery(
+            label="Generated images", 
+            show_label=False, 
+            elem_id="gallery"
+        )
+
+        with gr.Accordion("Advanced settings", open=False):
+            guidance_scale = gr.Slider(
+                label="Guidance Scale", 
+                minimum=1, 
+                maximum=20, 
+                value=9, 
+                step=0.1
+            )
+
+        gr.Examples(
+            examples=examples,
+            fn=infer,
+            inputs=[text, negative, guidance_scale],
+            outputs=[gallery],
+            cache_examples=True
+        )
+
+        text.submit(infer, inputs=[text, negative, guidance_scale], outputs=[gallery])
+        negative.submit(infer, inputs=[text, negative, guidance_scale], outputs=[gallery])
+        btn.click(infer, inputs=[text, negative, guidance_scale], outputs=[gallery])
+
+    gr.HTML(
+        """
+        <div style="text-align: center; margin-top: 20px;">
+            <p>Running on CPU - Please allow longer generation times</p>
+        </div>
+        """
+    )
+
+block.queue().launch(show_error=True)