Spaces:
Running
Running
Jose Benitez
commited on
Commit
•
025cc15
1
Parent(s):
c3fc33a
clean code
Browse files- gradio_app.py +52 -82
- routes.py +0 -11
- static/html/landing.html +16 -16
- 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 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
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 |
-
|
14 |
-
|
15 |
-
if not isinstance(lora_models, list):
|
16 |
raise ValueError("Expected loras_models to be a list of dictionaries.")
|
17 |
|
18 |
-
|
19 |
-
|
20 |
-
|
21 |
-
|
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 |
-
|
36 |
-
|
37 |
-
|
|
|
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 =
|
53 |
-
new_placeholder = f"
|
54 |
trigger_word = selected_lora["trigger_word"]
|
55 |
-
updated_text = f"####
|
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
|
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("
|
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("
|
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=
|
183 |
with gr.Column(elem_id="google-btn-container", elem_classes="google-btn-container svelte-vt1mxs gap"):
|
184 |
-
btn = gr.Button("
|
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(
|
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=
|
198 |
-
title = gr.HTML(
|
199 |
|
200 |
with gr.Column(elem_id="logout-btn-container"):
|
201 |
-
gr.Button("
|
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("
|
215 |
|
216 |
|
217 |
with gr.Tabs():
|
218 |
-
with gr.TabItem('
|
219 |
with gr.Row():
|
220 |
with gr.Column(scale=3):
|
221 |
prompt = gr.Textbox(label="Prompt",
|
222 |
lines=1,
|
223 |
-
placeholder="
|
224 |
-
info='
|
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("
|
234 |
selected_info = gr.Markdown("")
|
235 |
gallery = gr.Gallery(
|
236 |
-
[(item["image_url"], item["model_name"]) for item in
|
237 |
-
label="
|
238 |
allow_preview=False,
|
239 |
columns=3,
|
240 |
elem_id="gallery"
|
241 |
)
|
242 |
|
243 |
-
with gr.Accordion("
|
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("
|
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("
|
283 |
-
gr.Markdown("#
|
284 |
-
gr.Markdown("
|
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="
|
289 |
-
trigger_word = gr.Textbox(label="
|
290 |
-
info="
|
291 |
)
|
292 |
-
train_button = gr.Button("
|
293 |
-
with gr.Accordion("
|
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 |
-
|
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 |
-
<!--
|
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>🚀
|
148 |
-
<p>
|
149 |
</div>
|
150 |
</section>
|
151 |
|
152 |
<section class="features">
|
153 |
-
<h2>🌟
|
154 |
<div class="feature-grid">
|
155 |
<div class="feature">
|
156 |
-
<h3>
|
157 |
-
<p>
|
158 |
</div>
|
159 |
<div class="feature">
|
160 |
-
<h3>
|
161 |
-
<p>
|
162 |
</div>
|
163 |
<div class="feature">
|
164 |
-
<h3>
|
165 |
-
<p>
|
166 |
</div>
|
167 |
</div>
|
168 |
</section>
|
169 |
|
170 |
<section class="testimonials">
|
171 |
<div class="container">
|
172 |
-
<h2>💬
|
173 |
<div class="testimonial">
|
174 |
-
<p>"ObtuAI
|
175 |
-
<p><strong>- Ana,
|
176 |
</div>
|
177 |
<div class="testimonial">
|
178 |
-
<p>"
|
179 |
-
<p><strong>- Carlos,
|
180 |
</div>
|
181 |
</div>
|
182 |
</section>
|
183 |
</div>
|
184 |
|
185 |
<footer>
|
186 |
-
<p>ObtuAI -
|
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 ""
|