LobsterQQQ commited on
Commit
f1c7b01
·
1 Parent(s): bc0242a

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +31 -0
app.py ADDED
@@ -0,0 +1,31 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from contextlib import nullcontext
2
+ import gradio as gr
3
+ import torch
4
+ from torch import autocast
5
+ from diffusers import StableDiffusionPipeline
6
+
7
+
8
+ device = "cuda" if torch.cuda.is_available() else "cpu"
9
+ context = autocast if device == "cuda" else nullcontext
10
+ dtype = torch.float16 if device == "cuda" else torch.float32
11
+
12
+ pipe = StableDiffusionPipeline.from_pretrained("ringhyacinth/nail-set-diffuser", torch_dtype=dtype)
13
+ pipe = pipe.to(device)
14
+
15
+
16
+ # Disable nsfw checker
17
+ disable_safety = True
18
+
19
+ if disable_safety:
20
+ def null_safety(images, **kwargs):
21
+ return images, False
22
+ pipe.safety_checker = null_safety
23
+
24
+
25
+ def infer(prompt, n_samples, steps, scale):
26
+
27
+ with context("cuda"):
28
+ images = pipe(n_samples*[prompt], guidance_scale=scale, num_inference_steps=steps).images
29
+
30
+ return images
31
+