File size: 1,547 Bytes
6001e3c
6bc9074
e845246
8ee496d
922fdb6
6bc9074
6001e3c
6bc9074
8ee496d
e845246
 
 
 
6bc9074
 
922fdb6
 
 
 
 
 
223ef25
 
5de7ece
8c7013a
922fdb6
 
8c7013a
922fdb6
 
 
 
8c7013a
 
 
 
 
 
 
6bc9074
 
 
8c7013a
 
 
 
 
 
 
 
 
6bc9074
 
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
import gradio as gr
import spaces
import torch
import os
from compel import Compel, ReturnedEmbeddingsType
from diffusers import DiffusionPipeline


model_name = os.environ.get('MODEL_NAME', 'UnfilteredAI/NSFW-gen-v2')
pipe = DiffusionPipeline.from_pretrained(
    model_name,
    torch_dtype=torch.float16
)
pipe.to('cuda')

compel = Compel(
  tokenizer=[pipe.tokenizer, pipe.tokenizer_2] ,
  text_encoder=[pipe.text_encoder, pipe.text_encoder_2],
  returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED,
  requires_pooled=[False, True]
)


@spaces.GPU(duration=120)
def generate(prompt, negative_prompt, num_inference_steps, guidance_scale, width, height, num_samples):
    embeds, pooled = compel(prompt)
    neg_embeds, neg_pooled = compel(negative_prompt)
    return pipe(
        prompt_embeds=embeds,
        pooled_prompt_embeds=pooled,
        negative_prompt_embeds=neg_embeds,
        negative_pooled_prompt_embeds=neg_pooled,
        num_inference_steps=num_inference_steps,
        guidance_scale=guidance_scale,
        width=width,
        height=height,
        num_images_per_prompt=num_samples
    ).images


gr.Interface(
    fn=generate,
    inputs=[
        gr.Text(label="Prompt"),
        gr.Text("", label="Negative Prompt"),
        gr.Number(7, label="Number inference steps"),
        gr.Number(3, label="Guidance scale"),
        gr.Number(512, label="Width"),
        gr.Number(512, label="Height"),
        gr.Number(1, label="# images"),
    ],
    outputs=gr.Gallery(),
).launch()