trashchenkov commited on
Commit
759b0d2
·
verified ·
1 Parent(s): 4e3fce0

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +112 -94
app.py CHANGED
@@ -1,56 +1,74 @@
1
  import gradio as gr
2
  import numpy as np
3
- import random
4
-
5
- # import spaces #[uncomment to use ZeroGPU]
6
- from diffusers import DiffusionPipeline
7
  import torch
 
 
8
 
 
9
  device = "cuda" if torch.cuda.is_available() else "cpu"
10
- model_repo_id = "stabilityai/sdxl-turbo" # Replace to the model you would like to use
 
 
 
 
 
 
11
 
12
- if torch.cuda.is_available():
13
- torch_dtype = torch.float16
14
- else:
15
- torch_dtype = torch.float32
16
 
17
- pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
18
- pipe = pipe.to(device)
19
 
20
  MAX_SEED = np.iinfo(np.int32).max
21
  MAX_IMAGE_SIZE = 1024
22
 
23
-
24
- # @spaces.GPU #[uncomment to use ZeroGPU]
25
  def infer(
 
26
  prompt,
27
  negative_prompt,
28
  seed,
29
- randomize_seed,
30
  width,
31
  height,
32
  guidance_scale,
33
  num_inference_steps,
34
  progress=gr.Progress(track_tqdm=True),
35
  ):
36
- if randomize_seed:
37
- seed = random.randint(0, MAX_SEED)
38
-
39
- generator = torch.Generator().manual_seed(seed)
40
-
41
- image = pipe(
42
- prompt=prompt,
43
- negative_prompt=negative_prompt,
44
- guidance_scale=guidance_scale,
45
- num_inference_steps=num_inference_steps,
46
- width=width,
47
- height=height,
48
- generator=generator,
49
- ).images[0]
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
50
 
51
  return image, seed
52
 
53
-
54
  examples = [
55
  "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
56
  "An astronaut riding a green horse",
@@ -66,89 +84,89 @@ css = """
66
 
67
  with gr.Blocks(css=css) as demo:
68
  with gr.Column(elem_id="col-container"):
69
- gr.Markdown(" # Text-to-Image Gradio Template")
70
-
71
- with gr.Row():
72
- prompt = gr.Text(
73
- label="Prompt",
74
- show_label=False,
75
- max_lines=1,
76
- placeholder="Enter your prompt",
77
- container=False,
78
- )
79
-
80
- run_button = gr.Button("Run", scale=0, variant="primary")
81
-
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
82
  result = gr.Image(label="Result", show_label=False)
83
 
84
  with gr.Accordion("Advanced Settings", open=False):
85
- negative_prompt = gr.Text(
86
- label="Negative prompt",
87
- max_lines=1,
88
- placeholder="Enter a negative prompt",
89
- visible=False,
90
- )
91
-
92
- seed = gr.Slider(
93
- label="Seed",
94
- minimum=0,
95
- maximum=MAX_SEED,
96
- step=1,
97
- value=0,
98
- )
99
-
100
- randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
101
-
102
  with gr.Row():
103
  width = gr.Slider(
104
  label="Width",
105
  minimum=256,
106
  maximum=MAX_IMAGE_SIZE,
107
  step=32,
108
- value=1024, # Replace with defaults that work for your model
109
  )
110
-
111
  height = gr.Slider(
112
  label="Height",
113
  minimum=256,
114
  maximum=MAX_IMAGE_SIZE,
115
  step=32,
116
- value=1024, # Replace with defaults that work for your model
117
- )
118
-
119
- with gr.Row():
120
- guidance_scale = gr.Slider(
121
- label="Guidance scale",
122
- minimum=0.0,
123
- maximum=10.0,
124
- step=0.1,
125
- value=0.0, # Replace with defaults that work for your model
126
- )
127
-
128
- num_inference_steps = gr.Slider(
129
- label="Number of inference steps",
130
- minimum=1,
131
- maximum=50,
132
- step=1,
133
- value=2, # Replace with defaults that work for your model
134
  )
135
 
136
  gr.Examples(examples=examples, inputs=[prompt])
137
- gr.on(
138
- triggers=[run_button.click, prompt.submit],
139
- fn=infer,
140
- inputs=[
141
- prompt,
142
- negative_prompt,
143
- seed,
144
- randomize_seed,
145
- width,
146
- height,
147
- guidance_scale,
148
- num_inference_steps,
149
- ],
150
- outputs=[result, seed],
151
- )
152
 
153
  if __name__ == "__main__":
154
  demo.launch()
 
1
  import gradio as gr
2
  import numpy as np
 
 
 
 
3
  import torch
4
+ from diffusers import DiffusionPipeline
5
+ import re
6
 
7
+ # Устройство и параметры загрузки модели
8
  device = "cuda" if torch.cuda.is_available() else "cpu"
9
+ torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
10
+
11
+ # Регулярное выражение для проверки корректности модели
12
+ VALID_REPO_ID_REGEX = re.compile(r"^[a-zA-Z0-9._\-]+/[a-zA-Z0-9._\-]+$")
13
+
14
+ def is_valid_repo_id(repo_id):
15
+ return bool(VALID_REPO_ID_REGEX.match(repo_id)) and not repo_id.endswith(('-', '.'))
16
 
17
+ # Изначально загружаем модель по умолчанию
18
+ model_repo_id = "CompVis/stable-diffusion-v1-4"
19
+ pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype).to(device)
 
20
 
21
+ # Загружаем веса LoRA
22
+ pipe.load_lora_weights("AnastasiaSh/sticker-cat-lora3")
23
 
24
  MAX_SEED = np.iinfo(np.int32).max
25
  MAX_IMAGE_SIZE = 1024
26
 
 
 
27
  def infer(
28
+ model,
29
  prompt,
30
  negative_prompt,
31
  seed,
 
32
  width,
33
  height,
34
  guidance_scale,
35
  num_inference_steps,
36
  progress=gr.Progress(track_tqdm=True),
37
  ):
38
+ global model_repo_id, pipe
39
+
40
+ # Проверяем и загружаем новую модель, если она изменена
41
+ if model != model_repo_id:
42
+ if not is_valid_repo_id(model):
43
+ raise gr.Error(f"Некорректный идентификатор модели: '{model}'. Проверьте название.")
44
+ try:
45
+ new_pipe = DiffusionPipeline.from_pretrained(model, torch_dtype=torch_dtype).to(device)
46
+ # Повторно загружаем LoRA для новой модели
47
+ new_pipe.load_lora_weights("AnastasiaSh/sticker-cat-lora3")
48
+ pipe = new_pipe
49
+ model_repo_id = model
50
+ except Exception as e:
51
+ raise gr.Error(f"Не удалось загрузить модель '{model}'.\nОшибка: {e}")
52
+
53
+ # Генератор случайных чисел для детерминированности
54
+ generator = torch.Generator(device=device).manual_seed(seed)
55
+
56
+ # Генерация изображения
57
+ try:
58
+ image = pipe(
59
+ prompt=prompt,
60
+ negative_prompt=negative_prompt,
61
+ guidance_scale=guidance_scale,
62
+ num_inference_steps=num_inference_steps,
63
+ width=width,
64
+ height=height,
65
+ generator=generator,
66
+ ).images[0]
67
+ except Exception as e:
68
+ raise gr.Error(f"Ошибка при генерации изображения: {e}")
69
 
70
  return image, seed
71
 
 
72
  examples = [
73
  "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
74
  "An astronaut riding a green horse",
 
84
 
85
  with gr.Blocks(css=css) as demo:
86
  with gr.Column(elem_id="col-container"):
87
+ gr.Markdown("# Text-to-Image App")
88
+
89
+ model = gr.Textbox(
90
+ label="Model",
91
+ value="CompVis/stable-diffusion-v1-4", # Значение по умолчанию
92
+ interactive=True
93
+ )
94
+
95
+ prompt = gr.Text(
96
+ label="Prompt",
97
+ show_label=False,
98
+ max_lines=1,
99
+ placeholder="Enter your prompt",
100
+ container=False,
101
+ )
102
+
103
+ negative_prompt = gr.Text(
104
+ label="Negative prompt",
105
+ max_lines=1,
106
+ placeholder="Enter a negative prompt",
107
+ visible=True,
108
+ )
109
+
110
+ seed = gr.Slider(
111
+ label="Seed",
112
+ minimum=0,
113
+ maximum=MAX_SEED,
114
+ step=1,
115
+ value=42,
116
+ )
117
+
118
+ guidance_scale = gr.Slider(
119
+ label="Guidance scale",
120
+ minimum=0.0,
121
+ maximum=10.0,
122
+ step=0.1,
123
+ value=7.0,
124
+ )
125
+
126
+ num_inference_steps = gr.Slider(
127
+ label="Number of inference steps",
128
+ minimum=1,
129
+ maximum=50,
130
+ step=1,
131
+ value=20,
132
+ )
133
+
134
+ run_button = gr.Button("Run", variant="primary")
135
  result = gr.Image(label="Result", show_label=False)
136
 
137
  with gr.Accordion("Advanced Settings", open=False):
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
138
  with gr.Row():
139
  width = gr.Slider(
140
  label="Width",
141
  minimum=256,
142
  maximum=MAX_IMAGE_SIZE,
143
  step=32,
144
+ value=512,
145
  )
 
146
  height = gr.Slider(
147
  label="Height",
148
  minimum=256,
149
  maximum=MAX_IMAGE_SIZE,
150
  step=32,
151
+ value=512,
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
152
  )
153
 
154
  gr.Examples(examples=examples, inputs=[prompt])
155
+
156
+ run_button.click(
157
+ infer,
158
+ inputs=[
159
+ model,
160
+ prompt,
161
+ negative_prompt,
162
+ seed,
163
+ width,
164
+ height,
165
+ guidance_scale,
166
+ num_inference_steps,
167
+ ],
168
+ outputs=[result, seed],
169
+ )
170
 
171
  if __name__ == "__main__":
172
  demo.launch()