Pixelkram/webui
0
1import gradio as gr2import cv23import torch4import os5from imwatermark import WatermarkEncoder6import numpy as np7from PIL import Image8import re9from datasets import load_dataset10from diffusers import DiffusionPipeline, EulerDiscreteScheduler11 12from share_btn import community_icon_html, loading_icon_html, share_js13 14import huggingface_hub15#enter your own huggingface token here:16huggingface_hub.login(token="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")17 18REPO_ID = "stabilityai/stable-diffusion-2"19device = "cuda" if torch.cuda.is_available() else "cpu"20 21wm = "SDV2"22wm_encoder = WatermarkEncoder()23wm_encoder.set_watermark('bytes', wm.encode('utf-8'))24def put_watermark(img, wm_encoder=None):25 if wm_encoder is not None:26 img = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)27 img = wm_encoder.encode(img, 'dwtDct')28 img = Image.fromarray(img[:, :, ::-1])29 return img30 31repo_id = "stabilityai/stable-diffusion-2"32scheduler = EulerDiscreteScheduler.from_pretrained(repo_id, subfolder="scheduler", prediction_type="v_prediction")33pipe = DiffusionPipeline.from_pretrained(repo_id, torch_dtype=torch.float16, revision="fp16", scheduler=scheduler)34pipe = pipe.to(device)35#pipe.enable_xformers_memory_efficient_attention()36 37#If you have duplicated this Space or is running locally, you can remove this snippet38if "HUGGING_FACE_HUB_TOKEN" in os.environ:39 word_list_dataset = load_dataset("stabilityai/word-list", data_files="list.txt", use_auth_token=True)40 word_list = word_list_dataset["train"]['text']41 42def infer(prompt, samples, steps, scale, seed):43 #If you have duplicated this Space or is running locally, you can remove this snippet44 if "HUGGING_FACE_HUB_TOKEN" in os.environ:45 for filter in word_list:46 if re.search(rf"\b{filter}\b", prompt):47 raise gr.Error("Unsafe content found. Please try again with different prompts.")48 generator = torch.Generator(device=device).manual_seed(seed)49 images = pipe(prompt, width=768, height=768, num_inference_steps=steps, guidance_scale=scale, num_images_per_prompt=samples, generator=generator).images50 images_watermarked = []51 for image in images:52 image = put_watermark(image, wm_encoder)53 images_watermarked.append(image)54 return images_watermarked55 56css = """57 .gradio-container {58 font-family: 'IBM Plex Sans', sans-serif;59 }60 .gr-button {61 color: white;62 border-color: black;63 background: black;64 }65 input[type='range'] {66 accent-color: black;67 }68 .dark input[type='range'] {69 accent-color: #dfdfdf;70 }71 .container {72 max-width: 730px;73 margin: auto;74 padding-top: 1.5rem;75 }76 #gallery {77 min-height: 22rem;78 margin-bottom: 15px;79 margin-left: auto;80 margin-right: auto;81 border-bottom-right-radius: .5rem !important;82 border-bottom-left-radius: .5rem !important;83 }84 #gallery>div>.h-full {85 min-height: 20rem;86 }87 .details:hover {88 text-decoration: underline;89 }90 .gr-button {91 white-space: nowrap;92 }93 .gr-button:focus {94 border-color: rgb(147 197 253 / var(--tw-border-opacity));95 outline: none;96 box-shadow: var(--tw-ring-offset-shadow), var(--tw-ring-shadow), var(--tw-shadow, 0 0 #0000);97 --tw-border-opacity: 1;98 --tw-ring-offset-shadow: var(--tw-ring-inset) 0 0 0 var(--tw-ring-offset-width) var(--tw-ring-offset-color);99 --tw-ring-shadow: var(--tw-ring-inset) 0 0 0 calc(3px var(--tw-ring-offset-width)) var(--tw-ring-color);100 --tw-ring-color: rgb(191 219 254 / var(--tw-ring-opacity));101 --tw-ring-opacity: .5;102 }103 #advanced-btn {104 font-size: .7rem !important;105 line-height: 19px;106 margin-top: 12px;107 margin-bottom: 12px;108 padding: 2px 8px;109 border-radius: 14px !important;110 }111 #advanced-options {112 display: none;113 margin-bottom: 20px;114 }115 .footer {116 margin-bottom: 45px;117 margin-top: 35px;118 text-align: center;119 border-bottom: 1px solid #e5e5e5;120 }121 .footer>p {122 font-size: .8rem;123 display: inline-block;124 padding: 0 10px;125 transform: translateY(10px);126 background: white;127 }128 .dark .footer {129 border-color: #303030;130 }131 .dark .footer>p {132 background: #0b0f19;133 }134 .acknowledgments h4{135 margin: 1.25em 0 .25em 0;136 font-weight: bold;137 font-size: 115%;138 }139 .animate-spin {140 animation: spin 1s linear infinite;141 }142 @keyframes spin {143 from {144 transform: rotate(0deg);145 }146 to {147 transform: rotate(360deg);148 }149 }150 #share-btn-container {151 display: flex; padding-left: 0.5rem !important; padding-right: 0.5rem !important; background-color: #000000; justify-content: center; align-items: center; border-radius: 9999px !important; width: 13rem;152 margin-top: 10px;153 margin-left: auto;154 }155 #share-btn {156 all: initial; color: #ffffff;font-weight: 600; cursor:pointer; font-family: 'IBM Plex Sans', sans-serif; margin-left: 0.5rem !important; padding-top: 0.25rem !important; padding-bottom: 0.25rem !important;right:0;157 }158 #share-btn * {159 all: unset;160 }161 #share-btn-container div:nth-child(-n+2){162 width: auto !important;163 min-height: 0px !important;164 }165 #share-btn-container .wrap {166 display: none !important;167 }168 .gr-form{169 flex: 1 1 50%; border-top-right-radius: 0; border-bottom-right-radius: 0;170 }171 #prompt-container{172 gap: 0;173 }174 #component-9{margin-top: -19px}175 .image_duplication{position: absolute; width: 100px; left: 50px}176"""177 178block = gr.Blocks(css=css)179 180examples = [181 [182 'A high tech solarpunk utopia in the Amazon rainforest',183 4,184 25,185 9,186 1024,187 ],188 [189 'A pikachu fine dining with a view to the Eiffel Tower',190 4,191 25,192 9,193 1024,194 ],195 [196 'A mecha robot in a favela in expressionist style',197 4,198 25,199 9,200 1024,201 ],202 [203 'an insect robot preparing a delicious meal',204 4,205 25,206 9,207 1024,208 ],209 [210 "A small cabin on top of a snowy mountain in the style of Disney, artstation",211 4,212 25,213 9,214 1024,215 ],216]217 218with block:219 gr.HTML(220 """221 <div style="text-align: center; margin: 0 auto;">222 <div223 style="224 display: inline-flex;225 align-items: center;226 gap: 0.8rem;227 font-size: 1.75rem;228 "229 >230 <svg231 width="0.65em"232 height="0.65em"233 viewBox="0 0 115 115"234 fill="none"235 xmlns="http://www.w3.org/2000/svg"236 >237 <rect width="23" height="23" fill="white"></rect>238 <rect y="69" width="23" height="23" fill="white"></rect>239 <rect x="23" width="23" height="23" fill="#AEAEAE"></rect>240 <rect x="23" y="69" width="23" height="23" fill="#AEAEAE"></rect>241 <rect x="46" width="23" height="23" fill="white"></rect>242 <rect x="46" y="69" width="23" height="23" fill="white"></rect>243 <rect x="69" width="23" height="23" fill="black"></rect>244 <rect x="69" y="69" width="23" height="23" fill="black"></rect>245 <rect x="92" width="23" height="23" fill="#D9D9D9"></rect>246 <rect x="92" y="69" width="23" height="23" fill="#AEAEAE"></rect>247 <rect x="115" y="46" width="23" height="23" fill="white"></rect>248 <rect x="115" y="115" width="23" height="23" fill="white"></rect>249 <rect x="115" y="69" width="23" height="23" fill="#D9D9D9"></rect>250 <rect x="92" y="46" width="23" height="23" fill="#AEAEAE"></rect>251 <rect x="92" y="115" width="23" height="23" fill="#AEAEAE"></rect>252 <rect x="92" y="69" width="23" height="23" fill="white"></rect>253 <rect x="69" y="46" width="23" height="23" fill="white"></rect>254 <rect x="69" y="115" width="23" height="23" fill="white"></rect>255 <rect x="69" y="69" width="23" height="23" fill="#D9D9D9"></rect>256 <rect x="46" y="46" width="23" height="23" fill="black"></rect>257 <rect x="46" y="115" width="23" height="23" fill="black"></rect>258 <rect x="46" y="69" width="23" height="23" fill="black"></rect>259 <rect x="23" y="46" width="23" height="23" fill="#D9D9D9"></rect>260 <rect x="23" y="115" width="23" height="23" fill="#AEAEAE"></rect>261 <rect x="23" y="69" width="23" height="23" fill="black"></rect>262 </svg>263 <h1 style="font-weight: 900; margin-bottom: 7px;">264 Stable Diffusion 2 Demo265 </h1>266 </div>267 <p style="margin-bottom: 10px; font-size: 94%; line-height: 23px;">268 Stable Diffusion 2 is the latest text-to-image model from StabilityAI. <a style="text-decoration: underline;" href="https://huggingface.co/spaces/stabilityai/stable-diffusion-1">Access Stable Diffusion 1 Space here</a><br>For faster generation and API269 access you can try270 <a271 href="http://beta.dreamstudio.ai/"272 style="text-decoration: underline;"273 target="_blank"274 >DreamStudio Beta</a275 >.</a>276 </p>277 </div>278 """279 )280 with gr.Group():281 with gr.Box():282 with gr.Row(elem_id="prompt-container").style(mobile_collapse=False, equal_height=True):283 text = gr.Textbox(284 label="Enter your prompt",285 show_label=False,286 max_lines=1,287 placeholder="Enter your prompt",288 elem_id="prompt-text-input",289 ).style(290 border=(True, False, True, True),291 rounded=(True, False, False, True),292 container=False,293 )294 btn = gr.Button("Generate image").style(295 margin=False,296 rounded=(False, True, True, False),297 full_width=False,298 )299 300 gallery = gr.Gallery(301 label="Generated images", show_label=False, elem_id="gallery"302 ).style(grid=[2], height="auto")303 304 305 306 with gr.Accordion("Custom options", open=False):307 samples = gr.Slider(label="Images", minimum=1, maximum=4, value=4, step=1)308 steps = gr.Slider(label="Steps", minimum=1, maximum=50, value=25, step=1)309 scale = gr.Slider(310 label="Guidance Scale", minimum=0, maximum=50, value=9, step=0.1311 )312 seed = gr.Slider(313 label="Seed",314 minimum=0,315 maximum=2147483647,316 step=1,317 randomize=True,318 )319 320 with gr.Group():321 with gr.Group(elem_id="share-btn-container"):322 community_icon = gr.HTML(community_icon_html)323 loading_icon = gr.HTML(loading_icon_html)324 share_button = gr.Button("Share to community", elem_id="share-btn")325 326 ex = gr.Examples(examples=examples, fn=infer, inputs=[text, samples, steps, scale, seed], outputs=[gallery], cache_examples=False)327 ex.dataset.headers = [""]328 329 text.submit(infer, inputs=[text, samples, steps, scale, seed], outputs=[gallery])330 btn.click(infer, inputs=[text, samples, steps, scale, seed], outputs=[gallery])331 332 share_button.click(333 None,334 [],335 [],336 _js=share_js,337 )338 gr.HTML(339 """340 <div class="footer">341 <p>Model by <a href="https://huggingface.co/stabilityai" style="text-decoration: underline;" target="_blank">Stability AI</a> - Gradio Demo by 🤗 Hugging Face using the <a href="https://github.com/huggingface/diffusers" style="text-decoration: underline;" target="_blank">🧨 diffusers library</a>342 </p>343 </div>344 <div class="acknowledgments">345 <p><h4>LICENSE</h4>346The model is licensed with a <a href="https://huggingface.co/stabilityai/stable-diffusion-2/blob/main/LICENSE-MODEL" style="text-decoration: underline;" target="_blank">CreativeML OpenRAIL++</a> license. The authors claim no rights on the outputs you generate, you are free to use them and are accountable for their use which must not go against the provisions set in this license. 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>347 <p><h4>Biases and content acknowledgment</h4>348Despite 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 scraped non-curated image-text-pairs from the internet (the exception being 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>349 </div>350 """351 )352 353block.queue(concurrency_count=1, max_size=50).launch(max_threads=150)