anzorq commited on
Commit
a838e2b
1 Parent(s): 97a799e

+ custom models, colab support

Browse files
Files changed (1) hide show
  1. app.py +38 -26
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[0]
 
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(model_name, prompt, neg_prompt, img, strength, guidance, steps, width, height, generator)
49
  else:
50
- return txt_to_img(model_name, prompt, neg_prompt, guidance, steps, width, height, generator)
51
 
52
- def txt_to_img(model_name, prompt, neg_prompt, guidance, steps, width, height, generator=None):
53
 
54
- global current_model
55
  global pipe
56
- if model_name != current_model.name:
57
- for model in models:
58
- if model.name == model_name:
59
- current_model = model
60
 
61
- pipe = StableDiffusionPipeline.from_pretrained(current_model.path, torch_dtype=torch.float16)
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(model, prompt, neg_prompt, img, strength, guidance, steps, width, height, generator):
79
 
80
- global current_model
81
  global pipe
82
- if model_name != current_model.name:
83
- for model in models:
84
- if model.name == model_name:
85
- current_model = model
86
 
87
- pipe = StableDiffusionImg2ImgPipeline.from_pretrained(current_model.path, torch_dtype=torch.float16)
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=models[0].name)
 
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[0].name, "jason bateman disassembling the demon core", 7.5, 50],
180
- [models[3].name, "portrait of dwayne johnson", 7.0, 75],
181
- [models[4].name, "portrait of a beautiful alyx vance half life", 10, 50],
182
- [models[5].name, "Aloy from Horizon: Zero Dawn, half body portrait, smooth, detailed armor, beautiful face, illustration", 7.0, 45],
183
- [models[4].name, "fantasy portrait painting, digital art", 4.0, 30],
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)