Spaces:
Sleeping
Sleeping
File size: 4,693 Bytes
504d46f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 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) |