Jose Benitez commited on
Commit
025cc15
1 Parent(s): c3fc33a

clean code

Browse files
Files changed (4) hide show
  1. gradio_app.py +52 -82
  2. routes.py +0 -11
  3. static/html/landing.html +16 -16
  4. utils/file_utils.py +6 -0
gradio_app.py CHANGED
@@ -1,40 +1,35 @@
1
- import gradio as gr
2
-
3
  import os
4
- import json
5
  import zipfile
6
  from pathlib import Path
7
 
8
- from database import get_user_credits, update_user_credits, get_lora_models_info, get_user_lora_models
 
 
 
 
 
 
 
 
9
  from services.image_generation import generate_image
10
  from services.train_lora import lora_pipeline
11
  from utils.image_utils import url_to_pil_image
 
12
 
13
- lora_models = get_lora_models_info()
14
-
15
- if not isinstance(lora_models, list):
16
  raise ValueError("Expected loras_models to be a list of dictionaries.")
17
 
18
- login_css_path = Path(__file__).parent / 'static/css/login.css'
19
- main_css_path = Path(__file__).parent / 'static/css/main.css'
20
- landing_html_path = Path(__file__).parent / 'static/html/landing.html'
21
- main_header_path = Path(__file__).parent / 'static/html/main_header.html'
22
-
23
- if login_css_path.is_file(): # Check if the file exists
24
- with login_css_path.open() as file:
25
- login_css = file.read()
26
-
27
- if main_css_path.is_file(): # Check if the file exists
28
- with main_css_path.open() as file:
29
- main_css = file.read()
30
-
31
- if landing_html_path.is_file():
32
- with landing_html_path.open() as file:
33
- landin_page = file.read()
34
 
35
- if main_header_path.is_file():
36
- with main_header_path.open() as file:
37
- main_header = file.read()
 
38
 
39
  def load_user_models(request: gr.Request):
40
  user = request.session.get('user')
@@ -49,10 +44,10 @@ def update_selection(evt: gr.SelectData, gallery_type: str, width, height):
49
  if gallery_type == "user":
50
  selected_lora = {"lora_name": "custom", "trigger_word": "custom"}
51
  else:
52
- selected_lora = lora_models[evt.index]
53
- new_placeholder = f"Ingresa un prompt para tu modelo {selected_lora['lora_name']}"
54
  trigger_word = selected_lora["trigger_word"]
55
- updated_text = f"#### Palabra clave: {trigger_word} ✨"
56
 
57
  if "aspect" in selected_lora:
58
  if selected_lora["aspect"] == "portrait":
@@ -64,14 +59,14 @@ def update_selection(evt: gr.SelectData, gallery_type: str, width, height):
64
 
65
  def compress_and_train(request: gr.Request, files, model_name, trigger_word, train_steps, lora_rank, batch_size, learning_rate):
66
  if not files:
67
- return "No hay imágenes. Sube algunas imágenes para poder entrenar."
68
 
69
  user = request.session.get('user')
70
 
71
  _, training_credits = get_user_credits(user['id'])
72
 
73
  if training_credits <= 0:
74
- raise gr.Error("Ya no tienes creditos disponibles. Compra para continuar.")
75
 
76
  if not user:
77
  raise gr.Error("User not authenticated. Please log in.")
@@ -111,7 +106,7 @@ def compress_and_train(request: gr.Request, files, model_name, trigger_word, tra
111
  user['training_credits'] = new_training_credits
112
  request.session['user'] = user
113
 
114
- return gr.Info("Tu modelo esta entrenando, En unos 20 minutos estará listo para que lo pruebes en 'Generación'."), new_training_credits
115
 
116
  def run_lora(request: gr.Request, prompt, cfg_scale, steps, selected_index, selected_gallery, width, height, lora_scale, progress=gr.Progress(track_tqdm=True)):
117
  user = request.session.get('user')
@@ -179,9 +174,9 @@ def greet(request: gr.Request):
179
  return f"{greeting}\n"
180
  return "OBTU AI. Please log in."
181
 
182
- with gr.Blocks(theme=gr.themes.Soft(), css=login_css) as login_demo:
183
  with gr.Column(elem_id="google-btn-container", elem_classes="google-btn-container svelte-vt1mxs gap"):
184
- btn = gr.Button("Iniciar Sesion con Google", elem_classes="login-with-google-btn")
185
  _js_redirect = """
186
  () => {
187
  url = '/login' + window.location.search;
@@ -189,16 +184,16 @@ with gr.Blocks(theme=gr.themes.Soft(), css=login_css) as login_demo:
189
  }
190
  """
191
  btn.click(None, js=_js_redirect)
192
- gr.HTML(landin_page)
193
 
194
 
195
  header = '<script src="https://cdn.lordicon.com/lordicon.js"></script>'
196
 
197
- with gr.Blocks(theme=gr.themes.Soft(), head=header, css=main_css) as main_demo:
198
- title = gr.HTML(main_header)
199
 
200
  with gr.Column(elem_id="logout-btn-container"):
201
- gr.Button("Salir", link="/logout", elem_id="logout_btn")
202
 
203
 
204
  greetings = gr.Markdown("Loading user information...")
@@ -211,17 +206,17 @@ with gr.Blocks(theme=gr.themes.Soft(), head=header, css=main_css) as main_demo:
211
  with gr.Column():
212
  train_credits_display = gr.Number(label="Training Credits", precision=0, interactive=False)
213
  with gr.Column():
214
- gr.Button("Comprar Creditos 💳", link="/buy_credits")
215
 
216
 
217
  with gr.Tabs():
218
- with gr.TabItem('Generacion'):
219
  with gr.Row():
220
  with gr.Column(scale=3):
221
  prompt = gr.Textbox(label="Prompt",
222
  lines=1,
223
- placeholder="Ingresa un prompt para empezar a crear",
224
- info='Algunos modelos publicos pueden demorar un poco más dependiendo de la disponibilidad que tengan en los servidores.')
225
  with gr.Column(scale=1, elem_id="gen_column"):
226
  generate_button = gr.Button("Generate", variant="primary", elem_id="gen_btn")
227
 
@@ -230,17 +225,17 @@ with gr.Blocks(theme=gr.themes.Soft(), head=header, css=main_css) as main_demo:
230
  result = gr.Image(label="Imagen Generada")
231
 
232
  with gr.Column(scale=3):
233
- with gr.Accordion("Modelos Publicos"):
234
  selected_info = gr.Markdown("")
235
  gallery = gr.Gallery(
236
- [(item["image_url"], item["model_name"]) for item in lora_models],
237
- label="Galeria de Modelos Publicos",
238
  allow_preview=False,
239
  columns=3,
240
  elem_id="gallery"
241
  )
242
 
243
- with gr.Accordion("Tus Modelos"):
244
  user_model_gallery = gr.Gallery(
245
  label="Galeria de Modelos",
246
  allow_preview=False,
@@ -250,7 +245,7 @@ with gr.Blocks(theme=gr.themes.Soft(), head=header, css=main_css) as main_demo:
250
 
251
  gallery_type = gr.State("Public")
252
 
253
- with gr.Accordion("Configuracion Avanzada", open=False):
254
  with gr.Row():
255
  cfg_scale = gr.Slider(label="CFG Scale", minimum=1, maximum=20, step=0.5, value=3.5)
256
  steps = gr.Slider(label="Steps", minimum=1, maximum=50, step=1, value=28)
@@ -279,54 +274,29 @@ with gr.Blocks(theme=gr.themes.Soft(), head=header, css=main_css) as main_demo:
279
  outputs=[result, generation_credits_display]
280
  )
281
 
282
- with gr.TabItem("Training"):
283
- gr.Markdown("# Entrena tu propio modelo 🧠")
284
- gr.Markdown("En esta seccion podes entrenar tu propio modelo a partir de tus imagenes.")
285
  with gr.Row():
286
  with gr.Column():
287
  train_dataset = gr.Gallery(columns=4, interactive=True, label="Tus Imagenes")
288
- model_name = gr.Textbox(label="Nombre del Modelo",)
289
- trigger_word = gr.Textbox(label="Palabra clave",
290
- info="Esta seria una palabra clave para luego indicar al modelo cuando debe usar estas nuevas capacidad es que le vamos a ensenar",
291
  )
292
- train_button = gr.Button("Comenzar Training")
293
- with gr.Accordion("Configuracion Avanzada", open=False):
294
  train_steps = gr.Slider(label="Training Steps", minimum=100, maximum=10000, step=100, value=1000)
295
  lora_rank = gr.Number(label='lora_rank', value=16)
296
  batch_size = gr.Number(label='batch_size', value=1)
297
  learning_rate = gr.Number(label='learning_rate', value=0.0004)
298
  training_status = gr.Textbox(label="Training Status")
299
-
300
- def fake_train(train_dataset, model_name, trigger_word, train_steps, lora_rank, batch_size, learning_rate):
301
- print(f'fake training for test')
302
- new_training_credits = 0
303
- if new_training_credits <= 0:
304
- raise gr.Error("Ya no tienes creditos disponibles. Compra para continuar.")
305
- return gr.Info("Tu modelo esta entrenando, En unos 20 minutos estará listo para que lo pruebes en 'Generación'."), new_training_credits
306
 
307
  train_button.click(
308
- #compress_and_train,
309
- fake_train,
310
  inputs=[train_dataset, model_name, trigger_word, train_steps, lora_rank, batch_size, learning_rate],
311
  outputs=[training_status,train_credits_display]
312
  )
313
 
314
-
315
- #main_demo.load(greet, None, title)
316
- #main_demo.load(greet, None, greetings)
317
- #main_demo.load((greet, display_credits), None, [greetings, generation_credits_display, train_credits_display])
318
  main_demo.load(load_user_models, None, user_model_gallery)
319
- main_demo.load(load_greet_and_credits, None, [greetings, generation_credits_display, train_credits_display])
320
-
321
-
322
-
323
- # TODO:
324
- '''
325
- - resolver mostrar bien los nombres de los modelos en la galeria
326
- - Training con creditos.
327
- - Stripe(?)
328
- - Mejorar boton de login/logout
329
- - Retoque landing page
330
- '''
331
-
332
-
 
 
 
1
  import os
 
2
  import zipfile
3
  from pathlib import Path
4
 
5
+ import gradio as gr
6
+
7
+ from database import (
8
+ get_user_credits,
9
+ update_user_credits,
10
+ get_lora_models_info,
11
+ get_user_lora_models
12
+ )
13
+
14
  from services.image_generation import generate_image
15
  from services.train_lora import lora_pipeline
16
  from utils.image_utils import url_to_pil_image
17
+ from utils.file_utils import load_file_content
18
 
19
+ LORA_MODELS = get_lora_models_info()
20
+ if not isinstance(LORA_MODELS, list):
 
21
  raise ValueError("Expected loras_models to be a list of dictionaries.")
22
 
23
+ BASE_DIR = Path(__file__).parent
24
+ LOGIN_CSS_PATH = BASE_DIR / 'static/css/login.css'
25
+ MAIN_CSS_PATH = BASE_DIR / 'static/css/main.css'
26
+ LANDING_HTML_PATH = BASE_DIR / 'static/html/landing.html'
27
+ MAIN_HEADER_PATH = BASE_DIR / 'static/html/main_header.html'
 
 
 
 
 
 
 
 
 
 
 
28
 
29
+ LOGIN_CSS = load_file_content(LOGIN_CSS_PATH)
30
+ MAIN_CSS = load_file_content(MAIN_CSS_PATH)
31
+ LANDING_PAGE = load_file_content(LANDING_HTML_PATH)
32
+ MAIN_HEADER = load_file_content(MAIN_HEADER_PATH)
33
 
34
  def load_user_models(request: gr.Request):
35
  user = request.session.get('user')
 
44
  if gallery_type == "user":
45
  selected_lora = {"lora_name": "custom", "trigger_word": "custom"}
46
  else:
47
+ selected_lora = LORA_MODELS[evt.index]
48
+ new_placeholder = f"Enter a prompt for {selected_lora['lora_name']}"
49
  trigger_word = selected_lora["trigger_word"]
50
+ updated_text = f"#### Trigger Word: {trigger_word} ✨"
51
 
52
  if "aspect" in selected_lora:
53
  if selected_lora["aspect"] == "portrait":
 
59
 
60
  def compress_and_train(request: gr.Request, files, model_name, trigger_word, train_steps, lora_rank, batch_size, learning_rate):
61
  if not files:
62
+ return "No Images. Please, upload some images to start training"
63
 
64
  user = request.session.get('user')
65
 
66
  _, training_credits = get_user_credits(user['id'])
67
 
68
  if training_credits <= 0:
69
+ raise gr.Error("You ran out of credtis. Please buy more to continue")
70
 
71
  if not user:
72
  raise gr.Error("User not authenticated. Please log in.")
 
106
  user['training_credits'] = new_training_credits
107
  request.session['user'] = user
108
 
109
+ return gr.Info("Your model is training. In about 20 minutes, it will be ready for you to test in 'Generation"), new_training_credits
110
 
111
  def run_lora(request: gr.Request, prompt, cfg_scale, steps, selected_index, selected_gallery, width, height, lora_scale, progress=gr.Progress(track_tqdm=True)):
112
  user = request.session.get('user')
 
174
  return f"{greeting}\n"
175
  return "OBTU AI. Please log in."
176
 
177
+ with gr.Blocks(theme=gr.themes.Soft(), css=LOGIN_CSS) as login_demo:
178
  with gr.Column(elem_id="google-btn-container", elem_classes="google-btn-container svelte-vt1mxs gap"):
179
+ btn = gr.Button("Sign In with Google", elem_classes="login-with-google-btn")
180
  _js_redirect = """
181
  () => {
182
  url = '/login' + window.location.search;
 
184
  }
185
  """
186
  btn.click(None, js=_js_redirect)
187
+ gr.HTML(LANDING_PAGE)
188
 
189
 
190
  header = '<script src="https://cdn.lordicon.com/lordicon.js"></script>'
191
 
192
+ with gr.Blocks(theme=gr.themes.Soft(), head=header, css=MAIN_CSS) as main_demo:
193
+ title = gr.HTML(MAIN_HEADER)
194
 
195
  with gr.Column(elem_id="logout-btn-container"):
196
+ gr.Button("Logout", link="/logout", elem_id="logout_btn")
197
 
198
 
199
  greetings = gr.Markdown("Loading user information...")
 
206
  with gr.Column():
207
  train_credits_display = gr.Number(label="Training Credits", precision=0, interactive=False)
208
  with gr.Column():
209
+ gr.Button("Buy Credits 💳", link="/buy_credits")
210
 
211
 
212
  with gr.Tabs():
213
+ with gr.TabItem('Create'):
214
  with gr.Row():
215
  with gr.Column(scale=3):
216
  prompt = gr.Textbox(label="Prompt",
217
  lines=1,
218
+ placeholder="Enter Your Prompt to start creating 📷",
219
+ info='Some public models may experience longer processing times due to server availability and queue management.')
220
  with gr.Column(scale=1, elem_id="gen_column"):
221
  generate_button = gr.Button("Generate", variant="primary", elem_id="gen_btn")
222
 
 
225
  result = gr.Image(label="Imagen Generada")
226
 
227
  with gr.Column(scale=3):
228
+ with gr.Accordion("Public Models"):
229
  selected_info = gr.Markdown("")
230
  gallery = gr.Gallery(
231
+ [(item["image_url"], item["model_name"]) for item in LORA_MODELS],
232
+ label="Public Models",
233
  allow_preview=False,
234
  columns=3,
235
  elem_id="gallery"
236
  )
237
 
238
+ with gr.Accordion("Your Models"):
239
  user_model_gallery = gr.Gallery(
240
  label="Galeria de Modelos",
241
  allow_preview=False,
 
245
 
246
  gallery_type = gr.State("Public")
247
 
248
+ with gr.Accordion("Advanced Settings", open=False):
249
  with gr.Row():
250
  cfg_scale = gr.Slider(label="CFG Scale", minimum=1, maximum=20, step=0.5, value=3.5)
251
  steps = gr.Slider(label="Steps", minimum=1, maximum=50, step=1, value=28)
 
274
  outputs=[result, generation_credits_display]
275
  )
276
 
277
+ with gr.TabItem("Train"):
278
+ gr.Markdown("# Train your own model 🧠")
279
+ gr.Markdown("In this section, you can train your own model using your images.")
280
  with gr.Row():
281
  with gr.Column():
282
  train_dataset = gr.Gallery(columns=4, interactive=True, label="Tus Imagenes")
283
+ model_name = gr.Textbox(label="Model Name",)
284
+ trigger_word = gr.Textbox(label="Trigger Word",
285
+ info="This will be a keyword to later instruct the model when to use these new capabilities we're going to teach it",
286
  )
287
+ train_button = gr.Button("Start Training")
288
+ with gr.Accordion("Advanced Settings", open=False):
289
  train_steps = gr.Slider(label="Training Steps", minimum=100, maximum=10000, step=100, value=1000)
290
  lora_rank = gr.Number(label='lora_rank', value=16)
291
  batch_size = gr.Number(label='batch_size', value=1)
292
  learning_rate = gr.Number(label='learning_rate', value=0.0004)
293
  training_status = gr.Textbox(label="Training Status")
 
 
 
 
 
 
 
294
 
295
  train_button.click(
296
+ compress_and_train,
 
297
  inputs=[train_dataset, model_name, trigger_word, train_steps, lora_rank, batch_size, learning_rate],
298
  outputs=[training_status,train_credits_display]
299
  )
300
 
 
 
 
 
301
  main_demo.load(load_user_models, None, user_model_gallery)
302
+ main_demo.load(load_greet_and_credits, None, [greetings, generation_credits_display, train_credits_display])
 
 
 
 
 
 
 
 
 
 
 
 
 
routes.py CHANGED
@@ -105,17 +105,6 @@ async def stripe_webhook(request: Request):
105
 
106
  return {"status": "success"}
107
 
108
- # @router.get("/success")
109
- # async def payment_success(request: Request):
110
- # print("Payment successful")
111
- # user = request.session.get('user')
112
- # print(user)
113
- # if user:
114
- # updated_user = get_user_by_id(user['id'])
115
- # if updated_user:
116
- # request.session['user'] = updated_user
117
- # return RedirectResponse(url='/gradio', status_code=303)
118
- # return RedirectResponse(url='/login', status_code=303)
119
 
120
  @router.get("/cancel")
121
  async def payment_cancel(request: Request):
 
105
 
106
  return {"status": "success"}
107
 
 
 
 
 
 
 
 
 
 
 
 
108
 
109
  @router.get("/cancel")
110
  async def payment_cancel(request: Request):
static/html/landing.html CHANGED
@@ -135,7 +135,7 @@
135
  <div class="header-content">
136
  <div class="logo">🎨 ObtuAI</div>
137
  <div id="google-btn-container">
138
- <!-- El botón será insertado aquí por Gradio -->
139
  </div>
140
  </div>
141
  </div>
@@ -144,46 +144,46 @@
144
  <div class="container">
145
  <section class="hero">
146
  <div class="hero-content">
147
- <h1>🚀 Bienvenido al Futuro de la Creación Visual</h1>
148
- <p>Crea imágenes con IA en segundos. ¡Escribe tu idea y mira cómo se convierte en arte!</p>
149
  </div>
150
  </section>
151
 
152
  <section class="features">
153
- <h2>🌟 Descubre el Poder de la Generación de Imágenes por IA</h2>
154
  <div class="feature-grid">
155
  <div class="feature">
156
- <h3>Personaliza</h3>
157
- <p>Alimenta tu modelo con tus propias imágenes y estilos.</p>
158
  </div>
159
  <div class="feature">
160
- <h3>Entrena</h3>
161
- <p>Nuestra IA aprende de tus preferencias.</p>
162
  </div>
163
  <div class="feature">
164
- <h3>Crea</h3>
165
- <p>Genera imágenes que reflejen tu visión única.</p>
166
  </div>
167
  </div>
168
  </section>
169
 
170
  <section class="testimonials">
171
  <div class="container">
172
- <h2>💬 Lo Que Dicen Nuestros Usuarios</h2>
173
  <div class="testimonial">
174
- <p>"ObtuAI ha revolucionado mi proceso creativo. ¡Ahora puedo visualizar mis ideas más locas en minutos!"</p>
175
- <p><strong>- Ana, Diseñadora Gráfica</strong></p>
176
  </div>
177
  <div class="testimonial">
178
- <p>"Entrenar mi propio modelo fue sorprendentemente fácil. Ahora hago fotografías mías y de mis clientes en segundos."</p>
179
- <p><strong>- Carlos, Fotógrafo Profesional</strong></p>
180
  </div>
181
  </div>
182
  </section>
183
  </div>
184
 
185
  <footer>
186
- <p>ObtuAI - Tus ideas locas en píxeles con AI.</p>
187
  </footer>
188
  </body>
189
  </html>
 
135
  <div class="header-content">
136
  <div class="logo">🎨 ObtuAI</div>
137
  <div id="google-btn-container">
138
+ <!-- The button will be inserted here by Gradio -->
139
  </div>
140
  </div>
141
  </div>
 
144
  <div class="container">
145
  <section class="hero">
146
  <div class="hero-content">
147
+ <h1>🚀 Welcome to the Future of Visual Creation</h1>
148
+ <p>Create AI-generated images in seconds. Write your idea and watch it turn into art!</p>
149
  </div>
150
  </section>
151
 
152
  <section class="features">
153
+ <h2>🌟 Discover the Power of AI Image Generation</h2>
154
  <div class="feature-grid">
155
  <div class="feature">
156
+ <h3>Customize</h3>
157
+ <p>Feed your model with your own images and styles.</p>
158
  </div>
159
  <div class="feature">
160
+ <h3>Train</h3>
161
+ <p>Our AI learns from your preferences.</p>
162
  </div>
163
  <div class="feature">
164
+ <h3>Create</h3>
165
+ <p>Generate images that reflect your unique vision.</p>
166
  </div>
167
  </div>
168
  </section>
169
 
170
  <section class="testimonials">
171
  <div class="container">
172
+ <h2>💬 What Our Users Say</h2>
173
  <div class="testimonial">
174
+ <p>"ObtuAI has revolutionized my creative process. Now I can visualize my wildest ideas in minutes!"</p>
175
+ <p><strong>- Ana, Graphic Designer</strong></p>
176
  </div>
177
  <div class="testimonial">
178
+ <p>"Training my own model was surprisingly easy. Now I create photos of myself and my clients in seconds."</p>
179
+ <p><strong>- Carlos, Professional Photographer</strong></p>
180
  </div>
181
  </div>
182
  </section>
183
  </div>
184
 
185
  <footer>
186
+ <p>ObtuAI - Your wild ideas in pixels with AI.</p>
187
  </footer>
188
  </body>
189
  </html>
utils/file_utils.py ADDED
@@ -0,0 +1,6 @@
 
 
 
 
 
 
 
1
+ # Load static files
2
+ def load_file_content(file_path):
3
+ if file_path.is_file():
4
+ with file_path.open() as file:
5
+ return file.read()
6
+ return ""