add hf lora models support

services/image_generation.py

@@ -2,6 +2,7 @@ import replicate
 from PIL import Image
 import requests
 from io import BytesIO
+from database import get_lora_models_info
 
 
 def generate_image(model_name, prompt, steps, cfg_scale, width, height, lora_scale, progress, use_default=False, trigger_word='hi'):
@@ -19,15 +20,7 @@ def generate_image(model_name, prompt, steps, cfg_scale, width, height, lora_sca
             }
         )
     else:
-        # check if the model has a version. the model is something like model='user/model-name:version' but sometimes we just got model='user/model-name' in this case, let get and add the model version
-        if ':' not in model_name:
-            model_version = replicate.models.get(model_name).latest_version.id
-            print(f"Model version: {model_version}")
-            model_name = f"{model_name}:{model_version}"
-            
-        img_url = replicate.run(
-            model_name,
-            input={
+        input = {
                 "model": "dev",
                 "steps": steps,
                 "prompt": prompt,
@@ -36,5 +29,23 @@ def generate_image(model_name, prompt, steps, cfg_scale, width, height, lora_sca
                 "aspect_ratio": "1:1",
                 "safety_tolerance": 2
             }
+
+        db_loras = get_lora_models_info()
+
+        for lora in db_loras:
+            if lora["lora_name"] == model_name:
+                if lora["hf_repo"]:
+                    input["hf_lora"] = lora["hf_repo"]
+                    model_name = "lucataco/flux-dev-lora:a22c463f11808638ad5e2ebd582e07a469031f48dd567366fb4c6fdab91d614d"
+                    
+        if ':' not in model_name:
+            model_version = replicate.models.get(model_name).latest_version.id
+            print(f"Model version: {model_version}")
+            model_name = f"{model_name}:{model_version}"
+            
+        img_url = replicate.run(
+            model_name,
+            input=input
         )
     return img_url
+