Earendel/Inkpunk-Diffusion

1from diffusers import StableDiffusionPipeline, StableDiffusionImg2ImgPipeline, DPMSolverMultistepScheduler2import gradio as gr3import torch4from PIL import Image5 6model_id = 'Envvi/Inkpunk-Diffusion'7prefix = ''8     9scheduler = DPMSolverMultistepScheduler(10    beta_start=0.00085,11    beta_end=0.012,12    beta_schedule="scaled_linear",13    num_train_timesteps=1000,14    trained_betas=None,15    predict_epsilon=True,16    thresholding=False,17    algorithm_type="dpmsolver++",18    solver_type="midpoint",19    lower_order_final=True,20)21 22pipe = StableDiffusionPipeline.from_pretrained(23  model_id,24  torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,25  scheduler=scheduler)26 27pipe_i2i = StableDiffusionImg2ImgPipeline.from_pretrained(28  model_id,29  torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,30  scheduler=scheduler)31 32if torch.cuda.is_available():33  pipe = pipe.to("cuda")34  pipe_i2i = pipe_i2i.to("cuda")35 36def error_str(error, title="Error"):37    return f"""#### {title}38            {error}"""  if error else ""39 40def inference(prompt, guidance, steps, width=512, height=512, seed=0, img=None, strength=0.5, neg_prompt="", auto_prefix=True):41 42  generator = torch.Generator('cuda').manual_seed(seed) if seed != 0 else None43  prompt = f"{prefix} {prompt}" if auto_prefix else prompt44 45  try:46    if img is not None:47      return img_to_img(prompt, neg_prompt, img, strength, guidance, steps, width, height, generator), None48    else:49      return txt_to_img(prompt, neg_prompt, guidance, steps, width, height, generator), None50  except Exception as e:51    return None, error_str(e)52 53def txt_to_img(prompt, neg_prompt, guidance, steps, width, height, generator):54 55    result = pipe(56      prompt,57      negative_prompt = neg_prompt,58      num_inference_steps = int(steps),59      guidance_scale = guidance,60      width = width,61      height = height,62      generator = generator)63    64    return replace_nsfw_images(result)65 66def img_to_img(prompt, neg_prompt, img, strength, guidance, steps, width, height, generator):67 68    ratio = min(height / img.height, width / img.width)69    img = img.resize((int(img.width * ratio), int(img.height * ratio)), Image.LANCZOS)70    result = pipe_i2i(71        prompt,72        negative_prompt = neg_prompt,73        init_image = img,74        num_inference_steps = int(steps),75        strength = strength,76        guidance_scale = guidance,77        width = width,78        height = height,79        generator = generator)80        81    return replace_nsfw_images(result)82 83def replace_nsfw_images(results):84 85    for i in range(len(results.images)):86      if results.nsfw_content_detected[i]:87        results.images[i] = Image.open("nsfw.png")88    return results.images[0]89 90css = """.main-div div{display:inline-flex;align-items:center;gap:.8rem;font-size:1.75rem}.main-div div h1{font-weight:900;margin-bottom:7px}.main-div p{margin-bottom:10px;font-size:94%}a{text-decoration:underline}.tabs{margin-top:0;margin-bottom:0}#gallery{min-height:20rem}91"""92with gr.Blocks(css=css) as demo:93    gr.HTML(94        f"""95            <div class="main-div">96              <div>97                <h1>Inkpunk Diffusion</h1>98              </div>99              <p>100               Demo for <a href="https://huggingface.co/Envvi/Inkpunk-Diffusion">Inkpunk Diffusion</a> Stable Diffusion model.<br>101        Add the following tokens to your prompts for the model to work properly: <b></b>.102              </p>103              <p>This demo is currently on cpu, to use it upgrade to gpu by going to settings after duplicating this space: <a style="display:inline-block" href="https://huggingface.co/spaces/akhaliq/Inkpunk-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>104              Running on <b>{"GPU 🔥" if torch.cuda.is_available() else "CPU 🥶"}</b>105            </div>106        """107    )108    with gr.Row():109        110        with gr.Column(scale=55):111          with gr.Group():112              with gr.Row():113                prompt = gr.Textbox(label="Prompt", show_label=False, max_lines=2,placeholder=f"{prefix} [your prompt]").style(container=False)114                generate = gr.Button(value="Generate").style(rounded=(False, True, True, False))115 116              image_out = gr.Image(height=512)117          error_output = gr.Markdown()118 119        with gr.Column(scale=45):120          with gr.Tab("Options"):121            with gr.Group():122              neg_prompt = gr.Textbox(label="Negative prompt", placeholder="What to exclude from the image")123              auto_prefix = gr.Checkbox(label="Prefix styling tokens automatically ()", value=True)124 125              with gr.Row():126                guidance = gr.Slider(label="Guidance scale", value=7.5, maximum=15)127                steps = gr.Slider(label="Steps", value=25, minimum=2, maximum=75, step=1)128 129              with gr.Row():130                width = gr.Slider(label="Width", value=512, minimum=64, maximum=1024, step=8)131                height = gr.Slider(label="Height", value=512, minimum=64, maximum=1024, step=8)132 133              seed = gr.Slider(0, 2147483647, label='Seed (0 = random)', value=0, step=1)134 135          with gr.Tab("Image to image"):136              with gr.Group():137                image = gr.Image(label="Image", height=256, tool="editor", type="pil")138                strength = gr.Slider(label="Transformation strength", minimum=0, maximum=1, step=0.01, value=0.5)139 140    auto_prefix.change(lambda x: gr.update(placeholder=f"{prefix} [your prompt]" if x else "[Your prompt]"), inputs=auto_prefix, outputs=prompt, queue=False)141 142    inputs = [prompt, guidance, steps, width, height, seed, image, strength, neg_prompt, auto_prefix]143    outputs = [image_out, error_output]144    prompt.submit(inference, inputs=inputs, outputs=outputs)145    generate.click(inference, inputs=inputs, outputs=outputs)146 147    gr.HTML("""148    <div style="border-top: 1px solid #303030;">149      <br>150      <p>This space was created using <a href="https://huggingface.co/spaces/anzorq/sd-space-creator">SD Space Creator</a>.</p>151    </div>152    """)153 154demo.queue(concurrency_count=1)155demo.launch()156