Add ability to choose SD Model (on DreamStudio)

benrugg · Nov 30, 2022 · d0ebea2 · d0ebea2
1 parent e83a0a1
commit d0ebea2
Show file tree

Hide file tree

Showing 8 changed files with 109 additions and 63 deletions.
diff --git a/__init__.py b/__init__.py
@@ -2,7 +2,7 @@
     "name": "AI Render - Stable Diffusion in Blender",
     "description": "Create amazing images using Stable Diffusion AI",
     "author": "Ben Rugg",
-    "version": (0, 6, 3),
+    "version": (0, 6, 4),
     "blender": (3, 0, 0),
     "location": "Render Properties > AI Render",
     "warning": "",

diff --git a/config.py b/config.py
@@ -7,6 +7,7 @@
 ADDON_DOWNLOAD_URL = "https://airender.gumroad.com/l/ai-render"
 STABILITY_API_URL = "https://api.stability.ai/v1alpha/generation/"
 DREAM_STUDIO_URL = "https://beta.dreamstudio.ai/membership?tab=apiKeys"
+STABLE_HORDE_API_URL = "https://stablehorde.net/api/v2/generate/sync"
 STABLE_HORDE_URL = "https://stablehorde.net/"
 VIDEO_TUTORIAL_URL = "https://www.youtube.com/watch?v=tmyln5bwnO8"
 HELP_WITH_TIMEOUTS_URL = "https://github.com/benrugg/AI-Render/wiki/FAQ#%EF%B8%8F-ai-render-keeps-timing-out"

diff --git a/operators.py b/operators.py
@@ -498,7 +498,7 @@ def send_to_api(scene, prompts=None):
     }
 
     # send to whichever API we're using
-    output_file = utils.get_active_backend().send_to_api(params, img_file, after_output_filename_prefix)
+    output_file = utils.get_active_backend().send_to_api(params, img_file, after_output_filename_prefix, props.sd_model)
 
     # if we got a successful image created, handle it
     if output_file:

diff --git a/properties.py b/properties.py
@@ -78,6 +78,15 @@ class AIRProperties(bpy.types.PropertyGroup):
         max=150,
         description="How long to process the image. Values in the range of 25-50 generally work well. Higher values take longer (and use more credits) and won't necessarily improve results",
     )
+    sd_model: bpy.props.EnumProperty(
+        name="Stable Diffusion Model",
+        default=20,
+        items=[
+            ('v1-5', 'SD 1.5', '', 20),
+            ('v2-0', 'SD 2.0', '', 30),
+        ],
+        description="The Stable Diffusion model to use. 2.0 is more accurate with some types of images, and prompts differently from earlier versions. 1.5 is better for using artist names and art styles in prompts",
+    )
     sampler: bpy.props.EnumProperty(
         name="Sampler",
         default=120, # maps to DPM++ 2M, which is a good, fast sampler

diff --git a/sd_backends/automatic1111_api.py b/sd_backends/automatic1111_api.py
@@ -12,7 +12,7 @@
 
 # CORE FUNCTIONS:
 
-def send_to_api(params, img_file, filename_prefix):
+def send_to_api(params, img_file, filename_prefix, sd_model):
 
     # map the generic params to the specific ones for the Automatic1111 API
     map_params(params)
@@ -106,6 +106,15 @@ def handle_api_error(response):
         return operators.handle_error("An error occurred in the Automatic1111 Stable Diffusion server. Check the server logs for more info.")
 
 
+# PRIVATE SUPPORT FUNCTIONS:
+
+def map_params(params):
+    params["denoising_strength"] = round(1 - params["image_similarity"], 2)
+    params["sampler_index"] = params["sampler"]
+
+
+# PUBLIC SUPPORT FUNCTIONS:
+
 def get_samplers():
     # NOTE: Keep the number values (fourth item in the tuples) in sync with DreamStudio's
     # values (in stability_api.py). These act like an internal unique ID for Blender
@@ -140,12 +149,9 @@ def supports_negative_prompts():
     return True
 
 
-def max_image_size():
-    return 1024 * 1024
-
+def supports_choosing_model():
+    return False
 
-# SUPPORT FUNCTIONS:
 
-def map_params(params):
-    params["denoising_strength"] = round(1 - params["image_similarity"], 2)
-    params["sampler_index"] = params["sampler"]
+def max_image_size():
+    return 1024 * 1024
diff --git a/sd_backends/stability_api.py b/sd_backends/stability_api.py
@@ -10,7 +10,7 @@
 
 # CORE FUNCTIONS:
 
-def send_to_api(params, img_file, filename_prefix):
+def send_to_api(params, img_file, filename_prefix, sd_model):
 
     # map the generic params to the specific ones for the Stability API
     map_params(params)
@@ -30,7 +30,14 @@ def send_to_api(params, img_file, filename_prefix):
     }
 
     # prepare the URL
-    engine = "stable-diffusion-v1-5" # TODO: make this a param
+    if sd_model == 'v2-0':
+        if params["width"] >= 768 and params["height"] >= 768:
+            engine = f"stable-diffusion-768-{sd_model}"
+        else:
+            engine = f"stable-diffusion-512-{sd_model}"
+    else:
+        engine = f"stable-diffusion-{sd_model}"
+
     api_url = f"{config.STABILITY_API_URL}{engine}/image-to-image"
 
     # send the API request
@@ -102,6 +109,38 @@ def handle_api_error(response):
         return operators.handle_error(error_message, error_key)
 
 
+# PRIVATE SUPPORT FUNCTIONS:
+
+def map_params(params):
+    params["step_schedule_start"] = round(1 - params["image_similarity"], 2)
+    params["sampler"] = params["sampler"].upper()
+    params["text_prompts"] = [
+        {"text": params["prompt"], "weight": 1},
+    ]
+    if params["negative_prompt"]:
+        params["text_prompts"].append({"text": params["negative_prompt"], "weight": -1})
+
+
+def parse_message_for_error(message):
+    if "\"Authorization\" is missing" in message:
+        return "Your DreamStudio API key is missing. Please enter it above.", "api_key"
+    elif "Incorrect API key" in message or "Unauthenticated" in message:
+        return f"Your DreamStudio API key is incorrect. Please find it on the DreamStudio website, and re-enter it above. [DreamStudio website]({config.DREAM_STUDIO_URL})", "api_key"
+    elif "body.width must be" in message or "body.height must be" in message:
+        return "Invalid width or height. They must be one of the following values: 512, 576, 640, 704, 768, 832, 896, 960, 1024.", "dimensions"
+    elif "body.sampler must be" in message:
+        return "Invalid sampler. Please choose a new Sampler under 'Advanced Options'.", "sampler"
+    elif "body.cfg_scale must be" in message:
+        return "Invalid prompt strength. 'Prompt Strength' must be in the range 0-35.", "prompt_strength"
+    elif "body.seed must be" in message:
+        return "Invalid seed value. Please choose a new 'Seed'.", "seed"
+    elif "body.steps must be" in message:
+        return "Invalid number of steps. 'Steps' must be in the range 10-150.", "steps"
+    return "", ""
+
+
+# PUBLIC SUPPORT FUNCTIONS:
+
 def get_samplers():
     # NOTE: Keep the number values (fourth item in the tuples) in sync with the other
     # backends, like Automatic1111. These act like an internal unique ID for Blender
@@ -135,35 +174,9 @@ def supports_negative_prompts():
     return True
 
 
-def max_image_size():
-    return 1024 * 1024
-
-
-# SUPPORT FUNCTIONS:
-
-def map_params(params):
-    params["step_schedule_start"] = round(1 - params["image_similarity"], 2)
-    params["sampler"] = params["sampler"].upper()
-    params["text_prompts"] = [
-        {"text": params["prompt"], "weight": 1},
-    ]
-    if params["negative_prompt"]:
-        params["text_prompts"].append({"text": params["negative_prompt"], "weight": -1})
+def supports_choosing_model():
+    return True
 
 
-def parse_message_for_error(message):
-    if "\"Authorization\" is missing" in message:
-        return "Your DreamStudio API key is missing. Please enter it above.", "api_key"
-    elif "Incorrect API key" in message or "Unauthenticated" in message:
-        return f"Your DreamStudio API key is incorrect. Please find it on the DreamStudio website, and re-enter it above. [DreamStudio website]({config.DREAM_STUDIO_URL})", "api_key"
-    elif "body.width must be" in message or "body.height must be" in message:
-        return "Invalid width or height. They must be one of the following values: 512, 576, 640, 704, 768, 832, 896, 960, 1024.", "dimensions"
-    elif "body.sampler must be" in message:
-        return "Invalid sampler. Please choose a new Sampler under 'Advanced Options'.", "sampler"
-    elif "body.cfg_scale must be" in message:
-        return "Invalid prompt strength. 'Prompt Strength' must be in the range 0-35.", "prompt_strength"
-    elif "body.seed must be" in message:
-        return "Invalid seed value. Please choose a new 'Seed'.", "seed"
-    elif "body.steps must be" in message:
-        return "Invalid number of steps. 'Steps' must be in the range 10-150.", "steps"
-    return "", ""
+def max_image_size():
+    return 1024 * 1024
diff --git a/sd_backends/stablehorde_api.py b/sd_backends/stablehorde_api.py
@@ -9,31 +9,16 @@
     utils,
 )
 
-API_URL = "https://stablehorde.net/api/v2/generate/sync"
-SAMPLER_MAP = {
-    "DDIM": "k_ddim",
-    "PLMS": ""
-}
 
 # CORE FUNCTIONS:
 
-def send_to_api(params, img_file, filename_prefix):
+def send_to_api(params, img_file, filename_prefix, sd_model):
 
     # map the generic params to the specific ones for the Stable Horde API
-    stablehorde_params = {
-        "prompt": params["prompt"],
-        # add a base 64 encoded image to the params
-        "source_image": base64.b64encode(img_file.read()).decode(),
-        "params": {
-            "cfg_scale": params["cfg_scale"],
-            "width": params["width"],
-            "height": params["height"],
-            "denoising_strength": round(1 - params["image_similarity"], 2),
-            "seed": str(params["seed"]),
-            "steps": params["steps"],
-            "sampler_name": params["sampler"],
-        }
-    }
+    stablehorde_params = map_params(params)
+
+    # add a base 64 encoded image to the params
+    stablehorde_params["source_image"] = base64.b64encode(img_file.read()).decode()
 
     # close the image file
     img_file.close()
@@ -53,7 +38,7 @@ def send_to_api(params, img_file, filename_prefix):
     start_time = time.monotonic()
     print("Sending to Stable Horde")
     try:
-        response = requests.post(API_URL, json=stablehorde_params, headers=headers, timeout=request_timeout())
+        response = requests.post(config.STABLE_HORDE_API_URL, json=stablehorde_params, headers=headers, timeout=request_timeout())
         img_file.close()
     except requests.exceptions.ReadTimeout:
         img_file.close()
@@ -111,6 +96,26 @@ def handle_api_success(response, filename_prefix):
 def handle_api_error(response):
     return operators.handle_error("The Stable Horde server returned an error: " + str(response.content))
 
+
+# PRIVATE SUPPORT FUNCTIONS:
+
+def map_params(params):
+    return {
+        "prompt": params["prompt"],
+        "params": {
+            "cfg_scale": params["cfg_scale"],
+            "width": params["width"],
+            "height": params["height"],
+            "denoising_strength": round(1 - params["image_similarity"], 2),
+            "seed": str(params["seed"]),
+            "steps": params["steps"],
+            "sampler_name": params["sampler"],
+        }
+    }
+
+
+# PUBLIC SUPPORT FUNCTIONS:
+
 def get_samplers():
     # NOTE: Keep the number values (fourth item in the tuples) in sync with DreamStudio's
     # values (in stability_api.py). These act like an internal unique ID for Blender
@@ -146,5 +151,9 @@ def supports_negative_prompts():
     return False
 
 
+def supports_choosing_model():
+    return False
+
+
 def max_image_size():
     return 1024 * 1024
diff --git a/ui/ui_panels.py b/ui/ui_panels.py
@@ -241,6 +241,14 @@ def draw(self, context):
         sub = row.column()
         sub.prop(props, 'cfg_scale', text="", slider=False)
 
+        # SD Model
+        if utils.get_active_backend().supports_choosing_model():
+            row = layout.row()
+            sub = row.column()
+            sub.label(text="Model")
+            sub = row.column()
+            sub.prop(props, 'sd_model', text="")
+
         # Sampler
         row = layout.row()
         sub = row.column()