Spaces:
Running
on
Zero
Running
on
Zero
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() |