Merge pull request mix1009#154 from davidmartinrius/main

Implementation of Promptgen API

README.md

@@ -627,3 +627,62 @@ result1 = api.img2img(
 file_path = "face_swapped_image.png"
 result1.image.save(file_path)
 ```
+
+### Prompt generator API by [David Martin Rius](https://github.com/davidmartinrius/):
+
+
+This is an unofficial implementation to use the api of promptgen. 
+Before installing the extension you have to check if you already have an extension called Promptgen. If so, you need to uninstall it.
+Once uninstalled you can install it in two ways:
+
+#### 1. From the user interface
+![image](https://github.com/davidmartinrius/sdwebuiapi/assets/16558194/d879719f-bb9f-44a7-aef7-b893d117bbea)
+
+#### 2. From the command line
+
+cd stable-diffusion-webui/extensions
+
+git clone -b api-implementation https://github.com/davidmartinrius/stable-diffusion-webui-promptgen.git
+
+Once installed:
+```
+api = webuiapi.WebUIApi()
+
+result = api.list_prompt_gen_models()
+print("list of models")
+print(result)
+# you will get something like this:
+#['AUTOMATIC/promptgen-lexart', 'AUTOMATIC/promptgen-majinai-safe', 'AUTOMATIC/promptgen-majinai-unsafe']
+
+text = "a box"
+
+To create a prompt from a text:
+# by default model_name is "AUTOMATIC/promptgen-lexart"
+result = api.prompt_gen(text=text)
+
+# Using a different model
+result = api.prompt_gen(text=text, model_name="AUTOMATIC/promptgen-majinai-unsafe")
+
+#Complete usage
+result =  api.prompt_gen(
+        text=text, 
+        model_name="AUTOMATIC/promptgen-majinai-unsafe",
+        batch_count= 1,
+        batch_size=10,
+        min_length=20,
+        max_length=150,
+        num_beams=1,
+        temperature=1,
+        repetition_penalty=1,
+        length_preference=1,
+        sampling_mode="Top K",
+        top_k=12,
+        top_p=0.15
+    )
+
+# result is a list of prompts. You can iterate the list or just get the first result like this: result[0]
+
+```
+
+
+

webuiapi/webuiapi.py

@@ -413,10 +413,13 @@ def __init__(
         if baseurl is None:
             if use_https:
                 baseurl = f"https://{host}:{port}/sdapi/v1"
+                unofficial_baseurl = f"https://{host}:{port}"
             else:
                 baseurl = f"http://{host}:{port}/sdapi/v1"
+                unofficial_baseurl = f"http://{host}:{port}"
 
         self.baseurl = baseurl
+        self.unofficial_baseurl = unofficial_baseurl
         self.default_sampler = sampler
         self.default_steps = steps
 
@@ -939,7 +942,47 @@ def interrogate(self, image, model="clip"):
 
         response = self.session.post(url=f"{self.baseurl}/interrogate", json=payload)
         return self._to_api_result(response)
+
+    def list_prompt_gen_models(self):
+
+        print("CALLING list_prompt_gen_models", f"{self.unofficial_baseurl}/promptgen/list_models")
+        response = self.session.get(url=f"{self.unofficial_baseurl}/promptgen/list_models")
+        return response.json()['available_models']
+
+    def prompt_gen(self, 
+        model_name: str = "AUTOMATIC/promptgen-lexart",
+        batch_count: int = 1,
+        batch_size: int = 10,
+        text: str = "",
+        min_length: int = 20,
+        max_length: int = 150,
+        num_beams: int = 1,
+        temperature: float = 1,
+        repetition_penalty: float = 1,
+        length_preference: float = 1,
+        sampling_mode: str = "Top K",
+        top_k: float = 12,
+        top_p: float = 0.15,
+    ):
+        payload = {
+            "model_name": model_name,
+            "batch_count": batch_count,
+            "batch_size": batch_size,
+            "text": text,
+            "min_length": min_length,
+            "max_length": max_length,
+            "num_beams": num_beams,
+            "temperature": temperature,
+            "repetition_penalty": repetition_penalty,
+            "length_preference": length_preference,
+            "sampling_mode": sampling_mode,
+            "top_k": top_k,
+            "top_p": top_p
+        }
 
+        response = self.session.post(url=f"{self.unofficial_baseurl}/promptgen/generate", json=payload)
+        return response.json()['prompts']
+
     def interrupt(self):
         response = self.session.post(url=f"{self.baseurl}/interrupt")
         return response.json()