Spaces:
Running
on
Zero
Running
on
Zero
| from diffusers import ( | |
| StableDiffusionPipeline, | |
| DPMSolverMultistepScheduler, | |
| DiffusionPipeline, | |
| ) | |
| import gradio as gr | |
| import torch | |
| from PIL import Image | |
| import time | |
| import psutil | |
| import random | |
| from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker | |
| start_time = time.time() | |
| current_steps = 25 | |
| SAFETY_CHECKER = StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker", torch_dtype=torch.float16) | |
| UPSCALER = DiffusionPipeline.from_pretrained("stabilityai/sd-x2-latent-upscaler", torch_dtype=torch.float16) | |
| UPSCALER.to("cuda") | |
| UPSCALER.enable_xformers_memory_efficient_attention() | |
| class Model: | |
| def __init__(self, name, path=""): | |
| self.name = name | |
| self.path = path | |
| if path != "": | |
| self.pipe_t2i = StableDiffusionPipeline.from_pretrained( | |
| path, torch_dtype=torch.float16, safety_checker=SAFETY_CHECKER | |
| ) | |
| self.pipe_t2i.scheduler = DPMSolverMultistepScheduler.from_config( | |
| self.pipe_t2i.scheduler.config | |
| ) | |
| else: | |
| self.pipe_t2i = None | |
| models = [ | |
| Model("Stable Diffusion v1-4", "CompVis/stable-diffusion-v1-4"), | |
| ] | |
| MODELS = {m.name: m for m in models} | |
| device = "GPU 🔥" if torch.cuda.is_available() else "CPU 🥶" | |
| def error_str(error, title="Error"): | |
| return ( | |
| f"""#### {title} | |
| {error}""" | |
| if error | |
| else "" | |
| ) | |
| def inference( | |
| model_name, | |
| prompt, | |
| guidance, | |
| steps, | |
| width=512, | |
| height=512, | |
| seed=0, | |
| neg_prompt="", | |
| ): | |
| print(psutil.virtual_memory()) # print memory usage | |
| if seed == 0: | |
| seed = random.randint(0, 2147483647) | |
| generator = torch.Generator("cuda").manual_seed(seed) | |
| try: | |
| low_res_image, up_res_image = txt_to_img( | |
| model_name, | |
| prompt, | |
| neg_prompt, | |
| guidance, | |
| steps, | |
| width, | |
| height, | |
| generator, | |
| ) | |
| return low_res_image, up_res_image, f"Done. Seed: {seed}", | |
| except Exception as e: | |
| return None, None, error_str(e) | |
| def txt_to_img( | |
| model_name, | |
| prompt, | |
| neg_prompt, | |
| guidance, | |
| steps, | |
| width, | |
| height, | |
| generator, | |
| ): | |
| pipe = MODELS[model_name].pipe_t2i | |
| if torch.cuda.is_available(): | |
| pipe = pipe.to("cuda") | |
| pipe.enable_xformers_memory_efficient_attention() | |
| low_res_latents = pipe( | |
| prompt, | |
| negative_prompt=neg_prompt, | |
| num_inference_steps=int(steps), | |
| guidance_scale=guidance, | |
| width=width, | |
| height=height, | |
| generator=generator, | |
| output_type="latent", | |
| ).images | |
| with torch.no_grad(): | |
| low_res_image = pipe.decode_latents(low_res_latents) | |
| low_res_image = pipe.numpy_to_pil(low_res_image)[0] | |
| up_res_image = UPSCALER( | |
| prompt=prompt, | |
| negative_prompt=neg_prompt, | |
| image=low_res_latents, | |
| num_inference_steps=20, | |
| guidance_scale=0, | |
| generator=generator, | |
| ).images[0] | |
| pipe.to("cpu") | |
| torch.cuda.empty_cache() | |
| return low_res_image, up_res_image | |
| def replace_nsfw_images(results): | |
| for i in range(len(results.images)): | |
| if results.nsfw_content_detected[i]: | |
| results.images[i] = Image.open("nsfw.png") | |
| return results.images | |
| with gr.Blocks(css="style.css") as demo: | |
| gr.HTML( | |
| f""" | |
| <div class="finetuned-diffusion-div"> | |
| <div> | |
| <h1>Protogen Diffusion</h1> | |
| </div> | |
| <p> | |
| Demo for multiple fine-tuned Protogen Stable Diffusion models. | |
| </p> | |
| <p> | |
| Running on <b>{device}</b> | |
| </p> | |
| <p>You can also duplicate this space and upgrade to gpu by going to settings:<br> | |
| <a style="display:inline-block" href="https://huggingface.co/spaces/patrickvonplaten/finetuned_diffusion?duplicate=true"><img src="https://img.shields.io/badge/-Duplicate%20Space-blue?labelColor=white&style=flat&logo=data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAABAAAAAQCAYAAAAf8/9hAAAAAXNSR0IArs4c6QAAAP5JREFUOE+lk7FqAkEURY+ltunEgFXS2sZGIbXfEPdLlnxJyDdYB62sbbUKpLbVNhyYFzbrrA74YJlh9r079973psed0cvUD4A+4HoCjsA85X0Dfn/RBLBgBDxnQPfAEJgBY+A9gALA4tcbamSzS4xq4FOQAJgCDwV2CPKV8tZAJcAjMMkUe1vX+U+SMhfAJEHasQIWmXNN3abzDwHUrgcRGmYcgKe0bxrblHEB4E/pndMazNpSZGcsZdBlYJcEL9Afo75molJyM2FxmPgmgPqlWNLGfwZGG6UiyEvLzHYDmoPkDDiNm9JR9uboiONcBXrpY1qmgs21x1QwyZcpvxt9NS09PlsPAAAAAElFTkSuQmCC&logoWidth=14" alt="Duplicate Space"></a></p> | |
| </div> | |
| """ | |
| ) | |
| with gr.Row(): | |
| with gr.Column(scale=55): | |
| with gr.Group(): | |
| model_name = gr.Dropdown( | |
| label="Model", | |
| choices=[m.name for m in models], | |
| value=models[0].name, | |
| ) | |
| with gr.Box(visible=False) as custom_model_group: | |
| custom_model_path = gr.Textbox( | |
| label="Custom model path", | |
| placeholder="Path to model, e.g. darkstorm2150/Protogen_x3.4_Official_Release", | |
| interactive=True, | |
| ) | |
| gr.HTML( | |
| "<div><font size='2'>Custom models have to be downloaded first, so give it some time.</font></div>" | |
| ) | |
| with gr.Row(): | |
| prompt = gr.Textbox( | |
| label="Prompt", | |
| show_label=False, | |
| max_lines=2, | |
| placeholder="Enter prompt.", | |
| ).style(container=False) | |
| generate = gr.Button(value="Generate").style( | |
| rounded=(False, True, True, False) | |
| ) | |
| # image_out = gr.Image(height=512) | |
| low_res_image = gr.Gallery( | |
| label="512-pix image", show_label=False, elem_id="gallery" | |
| ).style(grid=[2], height="auto") | |
| up_res_image = gr.Gallery( | |
| label="1024-pix image", show_label=False, elem_id="gallery" | |
| ).style(grid=[2], height="auto") | |
| state_info = gr.Textbox(label="State", show_label=False, max_lines=2).style( | |
| container=False | |
| ) | |
| error_output = gr.Markdown() | |
| with gr.Column(scale=45): | |
| with gr.Tab("Options"): | |
| with gr.Group(): | |
| neg_prompt = gr.Textbox( | |
| label="Negative prompt", | |
| placeholder="What to exclude from the image", | |
| ) | |
| with gr.Row(): | |
| guidance = gr.Slider( | |
| label="Guidance scale", value=7.5, maximum=15 | |
| ) | |
| steps = gr.Slider( | |
| label="Steps", | |
| value=current_steps, | |
| minimum=2, | |
| maximum=75, | |
| step=1, | |
| ) | |
| with gr.Row(): | |
| width = gr.Slider( | |
| label="Width", value=512, minimum=64, maximum=1024, step=8 | |
| ) | |
| height = gr.Slider( | |
| label="Height", value=512, minimum=64, maximum=1024, step=8 | |
| ) | |
| seed = gr.Slider( | |
| 0, 2147483647, label="Seed (0 = random)", value=0, step=1 | |
| ) | |
| inputs = [ | |
| model_name, | |
| prompt, | |
| guidance, | |
| steps, | |
| width, | |
| height, | |
| seed, | |
| neg_prompt, | |
| ] | |
| outputs = [low_res_image, up_res_image, error_output] | |
| prompt.submit(inference, inputs=inputs, outputs=outputs) | |
| generate.click(inference, inputs=inputs, outputs=outputs) | |
| # ex = gr.Examples( | |
| # [ | |
| # [models[0].name, "portrait of a beautiful alyx vance half life", 10, 50, "canvas frame, ((disfigured)), ((bad art)), ((deformed)),((extra limbs)),((close up)),((b&w)), weird colors, blurry, (((duplicate))), ((morbid)), ((mutilated)), [out of frame], extra fingers, mutated hands, ((poorly drawn hands)), ((poorly drawn face)), (((mutation))), (((deformed))), blurry, ((bad anatomy)), (((bad proportions))), ((extra limbs)), cloned face, (((disfigured))), out of frame, ugly, extra limbs, (bad anatomy), gross proportions, (malformed limbs), ((missing arms)), ((missing legs)), (((extra arms))), (((extra legs))), mutated hands, (fused fingers), (too many fingers), (((long neck))), Photoshop, video game, ugly, tiling, poorly drawn hands, poorly drawn feet, poorly drawn face, out of frame, mutation, mutated, extra limbs, extra legs, extra arms, disfigured, deformed, cross-eye, body out of frame, blurry, bad art, bad anatomy"], | |
| # ], | |
| # inputs=[model_name, prompt, guidance, steps, neg_prompt], | |
| # outputs=outputs, | |
| # fn=inference, | |
| # cache_examples=False, | |
| # ) | |
| gr.HTML( | |
| """ | |
| <div style="border-top: 1px solid #303030;"> | |
| <br> | |
| <p>Models by <a href="https://huggingface.co/darkstorm2150">@darkstorm2150</a> and others. ❤️</p> | |
| <p>This space uses the <a href="https://github.com/LuChengTHU/dpm-solver">DPM-Solver++</a> sampler by <a href="https://arxiv.org/abs/2206.00927">Cheng Lu, et al.</a>.</p> | |
| <p>Space by: Darkstorm (Victor Espinoza)<br> | |
| <a href="https://www.instagram.com/officialvictorespinoza/">Instagram</a> | |
| </div> | |
| """ | |
| ) | |
| print(f"Space built in {time.time() - start_time:.2f} seconds") | |
| demo.queue(concurrency_count=1) | |
| demo.launch() | |