updates

database.py

@@ -35,11 +35,33 @@ def get_or_create_user(google_id, email, name, given_name, profile_picture):
         return user.data[0]
     
 def get_lora_models_info():
-    lora_models = supabase.table("lora_models").select("*").execute()
+    lora_models = supabase.table("lora_models").select("*").is_("user_id", None).execute()
+    
     return lora_models.data
 
 def get_user_by_id(user_id):
     user = supabase.table("users").select("*").eq("id", user_id).execute()
     if user.data:
         return user.data[0]
-    return None
+    return None
+
+def create_lora_models(user_id, replicate_repo_name, trigger_word, steps, lora_rank, batch_size, learning_rate, hf_repo_name, training_url):
+    # create a jsonb from trigger_word, train_steps, lora_rank, batch_size, learning_rate values
+    model_config = {
+        "train_steps": steps,
+        "lora_rank": lora_rank,
+        "batch_size": batch_size,
+        "learning_rate": learning_rate
+    }
+    result = supabase.table("lora_models").insert({
+        "user_id": user_id,
+        "trigger_word": trigger_word,
+        "lora_name": replicate_repo_name,
+        "hf_repo": hf_repo_name,
+        "configs": model_config,
+        "training_url": training_url
+    }).execute()
+    
+def get_user_lora_models(user_id):
+    user_models = supabase.table("lora_models").select("*").eq("user_id", user_id).execute()
+    return user_models.data

gradio_app.py

@@ -5,14 +5,13 @@ import json
 import zipfile
 from pathlib import Path
 
-from database import get_user_credits, update_user_credits, get_lora_models_info
+from database import get_user_credits, update_user_credits, get_lora_models_info, get_user_lora_models
 from services.image_generation import generate_image
 from services.train_lora import lora_pipeline
 from utils.image_utils import url_to_pil_image
 
 lora_models = get_lora_models_info()
 
-
 if not isinstance(lora_models, list):
     raise ValueError("Expected loras_models to be a list of dictionaries.")
 
@@ -37,8 +36,20 @@ if main_header_path.is_file():
     with main_header_path.open() as file:
         main_header = file.read()
 
-def update_selection(evt: gr.SelectData, width, height):
-    selected_lora = lora_models[evt.index]
+def load_user_models(request: gr.Request):
+    user = request.session.get('user')
+    print(user)
+    if user:
+        user_models = get_user_lora_models(user['id'])
+        if user_models:
+            return [(item.get("image_url", "assets/logo.jpg"), item["lora_name"]) for item in user_models]
+    return []
+
+def update_selection(evt: gr.SelectData, gallery_type: str, width, height):
+    if gallery_type == "user":
+        selected_lora = {"lora_name": "custom", "trigger_word": "custom"}
+    else:
+        selected_lora = lora_models[evt.index]
     new_placeholder = f"Ingresa un prompt para tu modelo {selected_lora['lora_name']}"
     trigger_word = selected_lora["trigger_word"]
     updated_text = f"#### Palabra clave: {trigger_word} ✨"
@@ -48,13 +59,18 @@ def update_selection(evt: gr.SelectData, width, height):
             width, height = 768, 1024
         elif selected_lora["aspect"] == "landscape":
             width, height = 1024, 768
-    
-    return gr.update(placeholder=new_placeholder), updated_text, evt.index, width, height
 
-def compress_and_train(files, model_name, trigger_word, train_steps, lora_rank, batch_size, learning_rate):
+    return gr.update(placeholder=new_placeholder), updated_text, evt.index, width, height, gallery_type
+
+def compress_and_train(request: gr.Request, files, model_name, trigger_word, train_steps, lora_rank, batch_size, learning_rate):
     if not files:
-        return "No images uploaded. Please upload images before training."
+        return "No hay imágenes. Sube algunas imágenes para poder entrenar."
+    
+    user = request.session.get('user')
+    if not user:
+        raise gr.Error("User not authenticated. Please log in.")
 
+    user_id = user['id']
     # Create a directory in the user's home folder
     output_dir = os.path.expanduser("~/gradio_training_data")
     os.makedirs(output_dir, exist_ok=True)
@@ -72,7 +88,8 @@ def compress_and_train(files, model_name, trigger_word, train_steps, lora_rank,
         
     print(f'[INFO] Procesando {trigger_word}')
     # Now call the train_lora function with the zip file path
-    result = lora_pipeline(zip_path, 
+    result = lora_pipeline(user_id,
+                            zip_path, 
                             model_name, 
                             trigger_word=trigger_word, 
                             steps=train_steps, 
@@ -81,19 +98,39 @@ def compress_and_train(files, model_name, trigger_word, train_steps, lora_rank,
                             autocaption=True, 
                             learning_rate=learning_rate)
     
-    return f"{result}\n\nZip file saved at: {zip_path}"
+    return gr.Info("Tu modelo esta entrenando, En unos 20 minutos estará listo para que lo pruebes en 'Generación'.")
             
-def run_lora(request: gr.Request, prompt, cfg_scale, steps, selected_index, randomize_seed, width, height, lora_scale, progress=gr.Progress(track_tqdm=True)):
+def run_lora(request: gr.Request, prompt, cfg_scale, steps, selected_index, selected_gallery, width, height, lora_scale, progress=gr.Progress(track_tqdm=True)):
     user = request.session.get('user')
     if not user:
         raise gr.Error("User not authenticated. Please log in.")
-    
+    lora_models = get_user_lora_models(user['id'])
+    print(f'Selected gallery: {selected_gallery}')
+    if selected_gallery == "user":
+        lora_models = get_user_lora_models(user['id'])
+        print('Using user models')
+    else:  # public
+        lora_models = get_lora_models_info()
+        print('Using public models')
+    print(f'Selected index: {selected_index}')
+    if selected_index is None:
+        selected_lora = None
+    else:
+        selected_lora = lora_models[selected_index]
+
     generation_credits, _ = get_user_credits(user['id'])
-    
+    if selected_lora:
+        print(f"Selected Lora: {selected_lora['lora_name']}")
+        model_name = selected_lora['lora_name']
+        use_default = False
+    else:
+        model_name = "black-forest-labs/flux-pro"
+        print(f"Using default Lora: {model_name}")
+        use_default = True
     if generation_credits <= 0:
         raise gr.Error("Ya no tienes creditos disponibles. Compra para continuar.")
     
-    image_url = generate_image(prompt, steps, cfg_scale, width, height, lora_scale, progress)
+    image_url = generate_image(model_name, prompt, steps, cfg_scale, width, height, lora_scale, progress, use_default)
     image = url_to_pil_image(image_url)
     
     # Update user's credits
@@ -193,7 +230,7 @@ with gr.Blocks(theme=gr.themes.Soft(), head=header, css=main_css) as main_demo:
                             columns=3,
                             elem_id="gallery"
                         )
-                    
+            gallery_type = gr.State("Public")
 
             with gr.Accordion("Configuracion Avanzada", open=False):
                 with gr.Row():
@@ -208,14 +245,19 @@ with gr.Blocks(theme=gr.themes.Soft(), head=header, css=main_css) as main_demo:
 
             gallery.select(
                 update_selection,
-                inputs=[width, height],
-                outputs=[prompt, selected_info, selected_index, width, height]
+                inputs=[gr.State("public"), width, height],
+                outputs=[prompt, selected_info, selected_index, width, height, gallery_type]
+            )
+            
+            user_model_gallery.select(
+                update_selection,
+                inputs=[gr.State("user"), width, height],
+                outputs=[prompt, selected_info, selected_index, width, height, gallery_type]
             )
-
             gr.on(
                 triggers=[generate_button.click, prompt.submit],
                 fn=run_lora,
-                inputs=[prompt, cfg_scale, steps, selected_index, randomize_seed, width, height, lora_scale],
+                inputs=[prompt, cfg_scale, steps, selected_index, gallery_type, width, height, lora_scale],
                 outputs=[result, generation_credits_display]
             )
 
@@ -237,8 +279,6 @@ with gr.Blocks(theme=gr.themes.Soft(), head=header, css=main_css) as main_demo:
                 learning_rate = gr.Number(label='learning_rate', value=0.0004)
                 training_status = gr.Textbox(label="Training Status")
 
-
-
             train_button.click(
                 compress_and_train,
                 inputs=[train_dataset, model_name, trigger_word, train_steps, lora_rank, batch_size, learning_rate],
@@ -249,14 +289,14 @@ with gr.Blocks(theme=gr.themes.Soft(), head=header, css=main_css) as main_demo:
         #main_demo.load(greet, None, title)
         #main_demo.load(greet, None, greetings)
         #main_demo.load((greet, display_credits), None, [greetings, generation_credits_display, train_credits_display])
+        main_demo.load(load_user_models, None, user_model_gallery)
         main_demo.load(load_greet_and_credits, None, [greetings, generation_credits_display, train_credits_display])
 
 
 
 # TODO:
 '''
-- Galeria Modelos Propios (si existe alguno del user, si no, mostrar un mensaje para entrenar)
-- Galeria Modelos Open Source (accordion)
+- resolver mostrar bien los nombres de los modelos en la galeria
 - Training con creditos.
 - Stripe(?)
 - Mejorar boton de login/logout

services/image_generation.py

@@ -3,19 +3,32 @@ from PIL import Image
 import requests
 from io import BytesIO
 
-#model_custom_test = "josebenitezg/flux-dev-ruth-estilo-1:c7ff81b58007c7cee3f69416e1e999192dafd8d1b1f269ea6cae137f04b34172"
-flux_pro = "black-forest-labs/flux-pro"
-def generate_image(prompt, steps, cfg_scale, width, height, lora_scale, progress, trigger_word='hi'):
+
+def generate_image(model_name, prompt, steps, cfg_scale, width, height, lora_scale, progress, use_default=False, trigger_word='hi'):
     print(f"Generating image for prompt: {prompt}")
-    img_url = replicate.run(
-        flux_pro,
-        input={
-            "steps": steps,
-            "prompt": prompt,
-            "guidance": cfg_scale,
-            "interval": 2,
-            "aspect_ratio": "1:1",
-            "safety_tolerance": 2
-        }
-    )
+    if use_default:
+        img_url = replicate.run(
+            "black-forest-labs/flux-pro",
+            input={
+                "steps": steps,
+                "prompt": prompt,
+                "guidance": cfg_scale,
+                "interval": 2,
+                "aspect_ratio": "1:1",
+                "safety_tolerance": 2
+            }
+        )
+    else:
+        img_url = replicate.run(
+            model_name,
+            input={
+                "model": "dev",
+                "steps": steps,
+                "prompt": prompt,
+                "guidance": cfg_scale,
+                "interval": 2,
+                "aspect_ratio": "1:1",
+                "safety_tolerance": 2
+            }
+        )
     return img_url

services/train_lora.py

@@ -1,20 +1,22 @@
 import replicate
 import os
 from huggingface_hub import create_repo
+from database import create_lora_models
 
 REPLICATE_OWNER = "josebenitezg"
 
-def lora_pipeline(zip_path, model_name, trigger_word="TOK", steps=1000, lora_rank=16, batch_size=1, autocaption=True, learning_rate=0.0004):
+def lora_pipeline(user_id, zip_path, model_name, trigger_word="TOK", steps=1000, lora_rank=16, batch_size=1, autocaption=True, learning_rate=0.0004):
     print(f'Creating dataset for {model_name}')
-    repo_name = f"joselobenitezg/flux-dev-{model_name}"
-    create_repo(repo_name, repo_type='model')
+    hf_repo_name = f"joselobenitezg/flux-dev-{model_name}"
+    replicate_repo_name = f"josebenitezg/flux-dev-{model_name}"
+    create_repo(hf_repo_name, repo_type='model')
 
     lora_name = f"flux-dev-{model_name}"
     
     model = replicate.models.create(
         owner=REPLICATE_OWNER,
         name=lora_name,
-        visibility="public",  # or "private" if you prefer
+        visibility="private",  # or "private" if you prefer
         hardware="gpu-t4",  # Replicate will override this for fine-tuned models
         description="A fine-tuned FLUX.1 model"
     )
@@ -37,10 +39,14 @@ def lora_pipeline(zip_path, model_name, trigger_word="TOK", steps=1000, lora_ran
             "trigger_word": trigger_word,
             "learning_rate": learning_rate,
             "hf_token": os.getenv('HF_TOKEN'),  # optional
-            "hf_repo_id": repo_name,  # optional
+            "hf_repo_id": hf_repo_name,  # optional
         },
         destination=f"{model.owner}/{model.name}"
     )
 
+    print(f"training: {training.keys()}")
     print(f"Training started: {training.status}")
     print(f"Training URL: https://replicate.com/p/{training.id}")
+    print(f"Creating model in Database")
+    training_url = f"https://replicate.com/p/{training.id}"
+    create_lora_models(user_id, replicate_repo_name, trigger_word, steps, lora_rank, batch_size, learning_rate, hf_repo_name, training_url)