Spaces:
Runtime error
Runtime error
+ custom models, colab support
Browse files
app.py
CHANGED
@@ -17,6 +17,7 @@ class Model:
|
|
17 |
self.prefix = prefix
|
18 |
|
19 |
models = [
|
|
|
20 |
Model("Arcane", "nitrosocke/Arcane-Diffusion", "arcane style "),
|
21 |
Model("Archer", "nitrosocke/archer-diffusion", "archer style "),
|
22 |
Model("Elden Ring", "nitrosocke/elden-ring-diffusion", "elden ring style "),
|
@@ -32,33 +33,42 @@ models = [
|
|
32 |
Model("Hergé Style", "sd-dreambooth-library/herge-style", "herge_style "),
|
33 |
]
|
34 |
|
35 |
-
current_model = models[
|
|
|
36 |
pipe = StableDiffusionPipeline.from_pretrained(current_model.path, torch_dtype=torch.float16)
|
37 |
if torch.cuda.is_available():
|
38 |
pipe = pipe.to("cuda")
|
39 |
|
40 |
-
|
41 |
device = "GPU 🔥" if torch.cuda.is_available() else "CPU 🥶"
|
42 |
|
|
|
|
|
|
|
|
|
|
|
43 |
def inference(model_name, prompt, guidance, steps, width=512, height=512, seed=0, img=None, strength=0.5, neg_prompt=""):
|
44 |
|
|
|
|
|
|
|
|
|
|
|
|
|
45 |
generator = torch.Generator('cuda').manual_seed(seed) if seed != 0 else None
|
46 |
-
|
47 |
if img is not None:
|
48 |
-
return img_to_img(
|
49 |
else:
|
50 |
-
return txt_to_img(
|
51 |
|
52 |
-
def txt_to_img(
|
53 |
|
54 |
-
global current_model
|
55 |
global pipe
|
56 |
-
|
57 |
-
|
58 |
-
|
59 |
-
current_model = model
|
60 |
|
61 |
-
pipe = StableDiffusionPipeline.from_pretrained(
|
62 |
if torch.cuda.is_available():
|
63 |
pipe = pipe.to("cuda")
|
64 |
|
@@ -75,16 +85,14 @@ def txt_to_img(model_name, prompt, neg_prompt, guidance, steps, width, height, g
|
|
75 |
image = results.images[0] if not results.nsfw_content_detected[0] else Image.open("nsfw.png")
|
76 |
return image
|
77 |
|
78 |
-
def img_to_img(
|
79 |
|
80 |
-
global current_model
|
81 |
global pipe
|
82 |
-
|
83 |
-
|
84 |
-
|
85 |
-
current_model = model
|
86 |
|
87 |
-
pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
|
88 |
|
89 |
if torch.cuda.is_available():
|
90 |
pipe = pipe.to("cuda")
|
@@ -152,7 +160,8 @@ with gr.Blocks(css=css) as demo:
|
|
152 |
with gr.Row():
|
153 |
|
154 |
with gr.Column():
|
155 |
-
model_name = gr.Dropdown(label="Model", choices=[m.name for m in models], value=
|
|
|
156 |
prompt = gr.Textbox(label="Prompt", placeholder="Style prefix is applied automatically")
|
157 |
run = gr.Button(value="Run")
|
158 |
|
@@ -170,17 +179,20 @@ with gr.Blocks(css=css) as demo:
|
|
170 |
|
171 |
with gr.Column():
|
172 |
image_out = gr.Image(height=512)
|
|
|
173 |
|
|
|
|
|
174 |
inputs = [model_name, prompt, guidance, steps, width, height, seed, image, strength, neg_prompt]
|
175 |
prompt.submit(inference, inputs=inputs, outputs=image_out, scroll_to_output=True)
|
176 |
run.click(inference, inputs=inputs, outputs=image_out, scroll_to_output=True)
|
177 |
|
178 |
gr.Examples([
|
179 |
-
[models[
|
180 |
-
[models[
|
181 |
-
[models[
|
182 |
-
[models[
|
183 |
-
[models[
|
184 |
], [model_name, prompt, guidance, steps, seed], image_out, inference, cache_examples=not is_colab and torch.cuda.is_available())
|
185 |
|
186 |
gr.Markdown('''
|
@@ -192,4 +204,4 @@ with gr.Blocks(css=css) as demo:
|
|
192 |
|
193 |
if not is_colab:
|
194 |
demo.queue(concurrency_count=4)
|
195 |
-
demo.launch(debug=is_colab)
|
|
|
17 |
self.prefix = prefix
|
18 |
|
19 |
models = [
|
20 |
+
Model("Custom model", "", ""),
|
21 |
Model("Arcane", "nitrosocke/Arcane-Diffusion", "arcane style "),
|
22 |
Model("Archer", "nitrosocke/archer-diffusion", "archer style "),
|
23 |
Model("Elden Ring", "nitrosocke/elden-ring-diffusion", "elden ring style "),
|
|
|
33 |
Model("Hergé Style", "sd-dreambooth-library/herge-style", "herge_style "),
|
34 |
]
|
35 |
|
36 |
+
current_model = models[1]
|
37 |
+
current_model_path = current_model.path
|
38 |
pipe = StableDiffusionPipeline.from_pretrained(current_model.path, torch_dtype=torch.float16)
|
39 |
if torch.cuda.is_available():
|
40 |
pipe = pipe.to("cuda")
|
41 |
|
|
|
42 |
device = "GPU 🔥" if torch.cuda.is_available() else "CPU 🥶"
|
43 |
|
44 |
+
def custom_model_changed(path):
|
45 |
+
models[0].path = path
|
46 |
+
current_model = models[0]
|
47 |
+
return models[0].path
|
48 |
+
|
49 |
def inference(model_name, prompt, guidance, steps, width=512, height=512, seed=0, img=None, strength=0.5, neg_prompt=""):
|
50 |
|
51 |
+
global current_model
|
52 |
+
for model in models:
|
53 |
+
if model.name == model_name:
|
54 |
+
current_model = model
|
55 |
+
model_path = current_model.path
|
56 |
+
|
57 |
generator = torch.Generator('cuda').manual_seed(seed) if seed != 0 else None
|
58 |
+
|
59 |
if img is not None:
|
60 |
+
return img_to_img(model_path, prompt, neg_prompt, img, strength, guidance, steps, width, height, generator)
|
61 |
else:
|
62 |
+
return txt_to_img(model_path, prompt, neg_prompt, guidance, steps, width, height, generator)
|
63 |
|
64 |
+
def txt_to_img(model_path, prompt, neg_prompt, guidance, steps, width, height, generator=None):
|
65 |
|
|
|
66 |
global pipe
|
67 |
+
global current_model_path
|
68 |
+
if model_path != current_model_path:
|
69 |
+
current_model_path = model_path
|
|
|
70 |
|
71 |
+
pipe = StableDiffusionPipeline.from_pretrained(current_model_path, torch_dtype=torch.float16)
|
72 |
if torch.cuda.is_available():
|
73 |
pipe = pipe.to("cuda")
|
74 |
|
|
|
85 |
image = results.images[0] if not results.nsfw_content_detected[0] else Image.open("nsfw.png")
|
86 |
return image
|
87 |
|
88 |
+
def img_to_img(model_path, prompt, neg_prompt, img, strength, guidance, steps, width, height, generator):
|
89 |
|
|
|
90 |
global pipe
|
91 |
+
global current_model_path
|
92 |
+
if model_path != current_model_path:
|
93 |
+
current_model_path = model_path
|
|
|
94 |
|
95 |
+
pipe = StableDiffusionImg2ImgPipeline.from_pretrained(current_model_path, torch_dtype=torch.float16)
|
96 |
|
97 |
if torch.cuda.is_available():
|
98 |
pipe = pipe.to("cuda")
|
|
|
160 |
with gr.Row():
|
161 |
|
162 |
with gr.Column():
|
163 |
+
model_name = gr.Dropdown(label="Model", choices=[m.name for m in models], value=current_model.name)
|
164 |
+
custom_model_path = gr.Textbox(label="Custom model path", placeholder="Path to model, e.g. nitrosocke/Arcane-Diffusion", visible=False, interactive=True)
|
165 |
prompt = gr.Textbox(label="Prompt", placeholder="Style prefix is applied automatically")
|
166 |
run = gr.Button(value="Run")
|
167 |
|
|
|
179 |
|
180 |
with gr.Column():
|
181 |
image_out = gr.Image(height=512)
|
182 |
+
log = gr.Textbox()
|
183 |
|
184 |
+
model_name.change(lambda x: gr.update(visible = x == models[0].name), inputs=model_name, outputs=custom_model_path)
|
185 |
+
custom_model_path.change(custom_model_changed, inputs=custom_model_path, outputs=log)
|
186 |
inputs = [model_name, prompt, guidance, steps, width, height, seed, image, strength, neg_prompt]
|
187 |
prompt.submit(inference, inputs=inputs, outputs=image_out, scroll_to_output=True)
|
188 |
run.click(inference, inputs=inputs, outputs=image_out, scroll_to_output=True)
|
189 |
|
190 |
gr.Examples([
|
191 |
+
[models[1].name, "jason bateman disassembling the demon core", 7.5, 50],
|
192 |
+
[models[4].name, "portrait of dwayne johnson", 7.0, 75],
|
193 |
+
[models[5].name, "portrait of a beautiful alyx vance half life", 10, 50],
|
194 |
+
[models[6].name, "Aloy from Horizon: Zero Dawn, half body portrait, smooth, detailed armor, beautiful face, illustration", 7.0, 45],
|
195 |
+
[models[5].name, "fantasy portrait painting, digital art", 4.0, 30],
|
196 |
], [model_name, prompt, guidance, steps, seed], image_out, inference, cache_examples=not is_colab and torch.cuda.is_available())
|
197 |
|
198 |
gr.Markdown('''
|
|
|
204 |
|
205 |
if not is_colab:
|
206 |
demo.queue(concurrency_count=4)
|
207 |
+
demo.launch(debug=is_colab, share=is_colab)
|