lora url support for replicate

fooocusapi/utils/lora_manager.py

@@ -0,0 +1,56 @@
+"""
+Manager loras from url
+
+@author: TechnikMax
+@github: https://github.com/TechnikMax
+"""
+import hashlib
+import os
+import requests
+
+
+def _hash_url(url):
+    """Generates a hash value for a given URL."""
+    return hashlib.md5(url.encode('utf-8')).hexdigest()
+
+
+class LoraManager:
+    """
+    Manager loras from url
+    """
+    def __init__(self):
+        self.cache_dir = "/models/loras/"
+
+    def _download_lora(self, url):
+        """
+        Downloads a LoRa from a URL and saves it in the cache.
+        """
+        url_hash = _hash_url(url)
+        filepath = os.path.join(self.cache_dir, f"{url_hash}.safetensors")
+        file_name = f"{url_hash}.safetensors"
+
+        if not os.path.exists(filepath):
+            print(f"start download for: {url}")
+
+            try:
+                response = requests.get(url, timeout=10, stream=True)
+                response.raise_for_status()
+                with open(filepath, 'wb') as f:
+                    for chunk in response.iter_content(chunk_size=8192):
+                        f.write(chunk)
+                print(f"Download successfully, saved as {file_name}")
+
+            except Exception as e:
+                raise Exception(f"error downloading {url}: {e}") from e
+
+        else:
+            print(f"LoRa already downloaded {url}")
+        return file_name
+
+    def check(self, urls):
+        """Manages the specified LoRAs: downloads missing ones and returns their file names."""
+        paths = []
+        for url in urls:
+            path = self._download_lora(url)
+            paths.append(path)
+        return paths

predict.py

@@ -10,6 +10,7 @@
 
 from PIL import Image
 from cog import BasePredictor, BaseModel, Input, Path
+from fooocusapi.utils.lora_manager import LoraManager
 from fooocusapi.utils.file_utils import output_dir
 from fooocusapi.models.common.task import GenerationFinishReason
 from fooocusapi.parameters import (
@@ -72,6 +73,12 @@ def predict(
         image_seed: int = Input(
             default=-1,
             description="Seed to generate image, -1 for random"),
+        use_default_loras: bool = Input(
+            default=True,
+            description="Use default LoRAs"),
+        loras_custom_urls: str = Input(
+            default="",
+            description="Custom LoRAs URLs in the format 'url,weight' provide multiple seperated by ; (example 'url1,0.3;url2,0.1')"),
         sharpness: float = Input(
             default=2.0,
             ge=0.0, le=30.0),
@@ -182,7 +189,23 @@ def predict(
 
         base_model_name = default_base_model_name
         refiner_model_name = default_refiner_model_name
-        loras = copy.copy(default_loras)
+
+        lora_manager = LoraManager()
+
+        # Use default loras if selected
+        loras = copy.copy(default_loras) if use_default_loras else []
+
+        # add custom user loras if provided
+        if loras_custom_urls:
+            urls = [url.strip() for url in loras_custom_urls.split(';')]
+
+            loras_with_weights = [url.split(',') for url in urls]
+
+            custom_lora_paths = lora_manager.check([lw[0] for lw in loras_with_weights])
+            custom_loras = [[path, float(lw[1]) if len(lw) > 1 else 1.0] for path, lw in
+                            zip(custom_lora_paths, loras_with_weights)]
+
+            loras.extend(custom_loras)
 
         style_selections_arr = []
         for s in style_selections.strip().split(','):