kappa2/stable-diffusion
0
1import gradio as gr2 3import torch4from torch import autocast5from diffusers import StableDiffusionPipeline6from datasets import load_dataset7from PIL import Image 8import re9 10model_id = "CompVis/stable-diffusion-v1-4"11device = "cuda"12 13pipe = StableDiffusionPipeline.from_pretrained(model_id, use_auth_token=False, revision="fp16", torch_dtype=torch.float16)14pipe = pipe.to(device)15word_list_dataset = load_dataset("stabilityai/word-list", data_files="list.txt", use_auth_token=False)16word_list = word_list_dataset["train"]['text']17 18def infer(prompt, samples, steps, scale, seed):19 for filter in word_list:20 if re.search(rf"\b{filter}\b", prompt):21 raise Exception("Unsafe content found. Please try again with different prompts.")22 23 generator = torch.Generator(device=device).manual_seed(seed)24 with autocast("cuda"):25 images_list = pipe(26 [prompt] * samples,27 num_inference_steps=steps,28 guidance_scale=scale,29 generator=generator,30 )31 images = []32 safe_image = Image.open(r"unsafe.png")33 for i, image in enumerate(images_list["sample"]):34 if(images_list["nsfw_content_detected"][i]):35 images.append(safe_image)36 else:37 images.append(image)38 return images39 40 41css = """42 .gradio-container {43 font-family: 'IBM Plex Sans', sans-serif;44 }45 .gr-button {46 color: white;47 border-color: black;48 background: black;49 }50 input[type='range'] {51 accent-color: black;52 }53 .dark input[type='range'] {54 accent-color: #dfdfdf;55 }56 .container {57 max-width: 1070px;58 margin: auto;59 padding-top: 2rem;60 }61 #gallery {62 min-height: 22rem;63 margin-bottom: 15px;64 border-bottom-right-radius: .5rem !important;65 border-bottom-left-radius: .5rem !important;66 }67 #gallery>div>.h-full {68 min-height: 20rem;69 }70 .details:hover {71 text-decoration: underline;72 }73 .gr-button {74 white-space: nowrap;75 }76 .gr-button:focus {77 border-color: rgb(147 197 253 / var(--tw-border-opacity));78 outline: none;79 box-shadow: var(--tw-ring-offset-shadow), var(--tw-ring-shadow), var(--tw-shadow, 0 0 #0000);80 --tw-border-opacity: 1;81 --tw-ring-offset-shadow: var(--tw-ring-inset) 0 0 0 var(--tw-ring-offset-width) var(--tw-ring-offset-color);82 --tw-ring-shadow: var(--tw-ring-inset) 0 0 0 calc(3px var(--tw-ring-offset-width)) var(--tw-ring-color);83 --tw-ring-color: rgb(191 219 254 / var(--tw-ring-opacity));84 --tw-ring-opacity: .5;85 }86 #advanced-btn {87 font-size: .7rem !important;88 line-height: 19px;89 margin-top: 24px;90 margin-bottom: 12px;91 padding: 2px 8px;92 border-radius: 14px !important;93 }94 #advanced-options {95 display: none;96 margin-bottom: 20px;97 }98 .footer {99 margin-bottom: 25px;100 text-align: center;101 border-bottom: 1px solid #e5e5e5;102 }103 .footer>p {104 font-size: .8rem;105 display: inline-block;106 padding: 0 10px;107 transform: translateY(10px);108 background: white;109 }110 .dark .footer {111 border-color: #303030;112 }113 .dark .footer>p {114 background: #0b0f19;115 }116 .acknowledgments h4{117 margin: 1.25em 0 .25em 0;118 font-weight: bold;119 font-size: 115%;120 }121"""122 123block = gr.Blocks(css=css)124 125examples = [126 [127 'A high tech solarpunk utopia in the Amazon rainforest',128 3,129 40,130 7.5,131 1024,132 ],133 [134 'A pikachu fine dining with a view to the Eiffel Tower',135 3,136 40,137 7,138 1024,139 ],140 [141 'A mecha robot in a favela in expressionist style',142 3,143 40,144 7,145 1024,146 ],147 [148 'an insect robot preparing a delicious meal',149 3,150 40,151 7,152 1024,153 ],154 [155 "A small cabin on top of a snowy mountain in the style of disney, arstation",156 3,157 40,158 7,159 1024,160 ],161]162 163with block:164 gr.HTML(165 """166 <div style="text-align: center;">167 <div style="display: inline-flex; align-items: center; gap: .8rem; font-size: 1.75rem;">168 <svg width="0.65em" height="0.65em" viewBox="0 0 115 115" fill="none" xmlns="http://www.w3.org/2000/svg">169 <rect width="23" height="23" fill="white"/>170 <rect y="69" width="23" height="23" fill="white"/>171 <rect x="23" width="23" height="23" fill="#AEAEAE"/>172 <rect x="23" y="69" width="23" height="23" fill="#AEAEAE"/>173 <rect x="46" width="23" height="23" fill="white"/>174 <rect x="46" y="69" width="23" height="23" fill="white"/>175 <rect x="69" width="23" height="23" fill="black"/>176 <rect x="69" y="69" width="23" height="23" fill="black"/>177 <rect x="92" width="23" height="23" fill="#D9D9D9"/>178 <rect x="92" y="69" width="23" height="23" fill="#AEAEAE"/>179 <rect x="115" y="46" width="23" height="23" fill="white"/>180 <rect x="115" y="115" width="23" height="23" fill="white"/>181 <rect x="115" y="69" width="23" height="23" fill="#D9D9D9"/>182 <rect x="92" y="46" width="23" height="23" fill="#AEAEAE"/>183 <rect x="92" y="115" width="23" height="23" fill="#AEAEAE"/>184 <rect x="92" y="69" width="23" height="23" fill="white"/>185 <rect x="69" y="46" width="23" height="23" fill="white"/>186 <rect x="69" y="115" width="23" height="23" fill="white"/>187 <rect x="69" y="69" width="23" height="23" fill="#D9D9D9"/>188 <rect x="46" y="46" width="23" height="23" fill="black"/>189 <rect x="46" y="115" width="23" height="23" fill="black"/>190 <rect x="46" y="69" width="23" height="23" fill="black"/>191 <rect x="23" y="46" width="23" height="23" fill="#D9D9D9"/>192 <rect x="23" y="115" width="23" height="23" fill="#AEAEAE"/>193 <rect x="23" y="69" width="23" height="23" fill="black"/>194 </svg>195 <h1 style="font-weight: 900;">Stable Diffusion Spaces</h1>196 </div>197 <p style="margin-bottom: 20px;">Stable Diffusion is a state of the art text-to-image model that generates images from a text description. For faster generation and forthcoming API access you can try <a href="http://beta.dreamstudio.ai/" style="text-decoration: underline;" target="_blank">DreamStudio Beta</a></p>198 </div>199 """200 )201 with gr.Group():202 with gr.Box():203 with gr.Row().style(mobile_collapse=False, equal_height=True):204 text = gr.Textbox(205 label="Enter your prompt",206 show_label=False,207 max_lines=1,208 placeholder="Enter your prompt",209 ).style(210 border=(True, False, True, True),211 rounded=(True, False, False, True),212 container=False,213 )214 btn = gr.Button("Generate image").style(215 margin=False,216 rounded=(False, True, True, False),217 )218 219 gallery = gr.Gallery(220 label="Generated images", show_label=False, elem_id="gallery"221 ).style(grid=[3], height="auto")222 223 advanced_button = gr.Button("Advanced options", elem_id="advanced-btn")224 225 with gr.Row(elem_id="advanced-options"):226 samples = gr.Slider(label="Images", minimum=1, maximum=3, value=3, step=1)227 steps = gr.Slider(label="Steps", minimum=1, maximum=50, value=40, step=1)228 scale = gr.Slider(229 label="Guidance Scale", minimum=0, maximum=50, value=7.5, step=0.1230 )231 seed = gr.Slider(232 label="Random seed",233 minimum=0,234 maximum=2147483647,235 step=1,236 randomize=True,237 )238 239 ex = gr.Examples(examples=examples, fn=infer, inputs=[text, samples, steps, scale, seed], outputs=gallery, cache_examples=True)240 ex.dataset.headers = [""]241 242 243 text.submit(infer, inputs=[text, samples, steps, scale, seed], outputs=gallery)244 btn.click(infer, inputs=[text, samples, steps, scale, seed], outputs=gallery)245 advanced_button.click(246 None,247 [],248 text,249 _js="""250 () => {251 const options = document.querySelector("body > gradio-app").querySelector("#advanced-options");252 options.style.display = ["none", ""].includes(options.style.display) ? "flex" : "none";253 }""",254 )255 gr.HTML(256 """257 <div class="footer">258 <p>Model by <a href="https://huggingface.co/CompVis" style="text-decoration: underline;" target="_blank">CompVis</a> and <a href="https://huggingface.co/stabilityai" style="text-decoration: underline;" target="_blank">Stability AI</a> - Demo by 🤗 Hugging Face259 </p>260 </div>261 <div class="acknowledgments">262 <p><h4>LICENSE</h4>263The model is licensed with an <a href="https://huggingface.co/spaces/CompVis/stable-diffusion-license" style="text-decoration: underline;" target="_blank">CreativeML Open RAIL-M</a> license. The license states that the outputs that you make fully belong to you, and you are liable when sharing it. The license forbids you from sharing any content that violates any laws, produce any harm to a person, disseminate any personal information that would be meant for harm, spread misinformation and target vulnerable groups. For the full list of restrictions please <a href="https://huggingface.co/spaces/CompVis/stable-diffusion-license" target="_blank" style="text-decoration: underline;" target="_blank">read the license</a></p>264 <p><h4>Biases and content acknowledgment</h4>265Despite how impressive being able to turn text into image is, beware to the fact that this model may output content that reinforces or exacerbates societal biases, as well as realistic faces, pornography and violence. The model was trained on the <a href="https://laion.ai/blog/laion-5b/" style="text-decoration: underline;" target="_blank">LAION-5B dataset</a>, which scrapped non-curated image-text-pairs from the internet (the exception being the the removal of illegal content) and is meant for research purposes. You can read more in the <a href="https://huggingface.co/CompVis/stable-diffusion-v1-4" style="text-decoration: underline;" target="_blank">model card</a></p>266 </div>267 """268 )269 270block.queue(max_size=40).launch()