├── README.md ├── LICENSE ├── LoadAndRun.py └── StableDiffusionWebUI.py /README.md: -------------------------------------------------------------------------------- 1 | # StablenderDiffusion 2 | General Repo for blender and stable-diffusion integration 3 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2022 shellworld 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /LoadAndRun.py: -------------------------------------------------------------------------------- 1 | ## 2 | ## 3 | ## Stablender Diffusion Load Script- For loading stability-sdk into blender 4 | ## version: 0.0.0 5 | ## author: Shellworld (twitter: @shellworld1; github.com/shellward) 6 | ## 7 | ## 8 | 9 | 10 | import pathlib 11 | import sys 12 | import os 13 | import uuid 14 | import random 15 | import io 16 | import logging 17 | import time 18 | import mimetypes 19 | import pip 20 | 21 | #install pip packages 22 | 23 | pip.main(['install', 'grpcio', '--user']) 24 | pip.main(['install', 'python-dotenv', '--user']) 25 | pip.main(['install', 'stability-sdk', '--user']) 26 | pip.main(['install', 'grpcio-tools', '--user']) 27 | pip.main(['install', 'PIL', '--user']) 28 | 29 | #REPLACE THIS WITH YOUR OWN PATH- you'll need to figure out which version of python blender is running, then locate the site-packages directory. 30 | #This is not necessarily attached to your local python installation- for example in windows 11 it's in your roaming app data/python/py version directory. 31 | #This might be a little different in linux/mac- but you are looking for the python-site-packages directory to attach to the syspath 32 | 33 | PACKAGES_PATH = "C:\\Users\\Henry\\AppData\\Roaming\\Python\\Python310\\site-packages" 34 | 35 | ## REPLACE THIS WITH YOUR OWN API KEY, which can be found at https://beta.dreamstudio.ai/membership 36 | STABILITY_API_KEY = "YOUR API KEY" 37 | #Set the default size for your requested image 38 | size=(512,512) 39 | 40 | 41 | 42 | #Insert paths 43 | packages_path= PACKAGES_PATH 44 | sys.path.insert(0, packages_path ) 45 | 46 | interfaces_path = f"{packages_path}\\stability_sdk" 47 | sys.path.insert(1,interfaces_path) 48 | 49 | generation_path = f"{interfaces_path}\\interfaces\\gooseai\\generation" 50 | sys.path.insert(2,generation_path) 51 | 52 | #Import the stability_sdk 53 | import grpc 54 | from argparse import ArgumentParser, Namespace 55 | from typing import Dict, Generator, List, Union, Any, Sequence, Tuple 56 | from dotenv import load_dotenv 57 | from google.protobuf.json_format import MessageToJson 58 | 59 | 60 | #This is all from the client.py file in the stability-sdk 61 | load_dotenv() 62 | 63 | thisPath = pathlib.Path(__file__).parent.resolve() 64 | genPath = thisPath / "interfaces/gooseai/generation" 65 | sys.path.append(str(genPath)) 66 | 67 | import stability_sdk.interfaces.gooseai.generation.generation_pb2 as generation 68 | import stability_sdk.interfaces.gooseai.generation.generation_pb2_grpc as generation_grpc 69 | 70 | logger = logging.getLogger(__name__) 71 | logger.setLevel(level=logging.INFO) 72 | 73 | algorithms: Dict[str, int] = { 74 | "ddim": generation.SAMPLER_DDIM, 75 | "plms": generation.SAMPLER_DDPM, 76 | "k_euler": generation.SAMPLER_K_EULER, 77 | "k_euler_ancestral": generation.SAMPLER_K_EULER_ANCESTRAL, 78 | "k_heun": generation.SAMPLER_K_HEUN, 79 | "k_dpm_2": generation.SAMPLER_K_DPM_2, 80 | "k_dpm_2_ancestral": generation.SAMPLER_K_DPM_2_ANCESTRAL, 81 | "k_lms": generation.SAMPLER_K_LMS, 82 | } 83 | 84 | 85 | def get_sampler_from_str(s: str) -> generation.DiffusionSampler: 86 | """ 87 | Convert a string to a DiffusionSampler enum. 88 | 89 | :param s: The string to convert. 90 | :return: The DiffusionSampler enum. 91 | """ 92 | algorithm_key = s.lower().strip() 93 | algorithm = algorithms.get(algorithm_key, None) 94 | if algorithm is None: 95 | raise ValueError(f"unknown sampler {s}") 96 | return algorithm 97 | 98 | 99 | def process_artifacts_from_answers( 100 | prefix: str, 101 | answers: Union[ 102 | Generator[generation.Answer, None, None], Sequence[generation.Answer] 103 | ], 104 | write: bool = True, 105 | verbose: bool = False, 106 | ) -> Generator[Tuple[str, generation.Artifact], None, None]: 107 | """ 108 | Process the Artifacts from the Answers. 109 | 110 | :param prefix: The prefix for the artifact filenames. 111 | :param answers: The Answers to process. 112 | :param write: Whether to write the artifacts to disk. 113 | :param verbose: Whether to print the artifact filenames. 114 | :return: A Generator of tuples of artifact filenames and Artifacts, intended 115 | for passthrough. 116 | """ 117 | idx = 0 118 | for resp in answers: 119 | for artifact in resp.artifacts: 120 | artifact_p = f"{prefix}-{resp.request_id}-{resp.answer_id}-{idx}" 121 | if artifact.type == generation.ARTIFACT_IMAGE: 122 | ext = mimetypes.guess_extension(artifact.mime) 123 | contents = artifact.binary 124 | elif artifact.type == generation.ARTIFACT_CLASSIFICATIONS: 125 | ext = ".pb.json" 126 | contents = MessageToJson(artifact.classifier).encode("utf-8") 127 | elif artifact.type == generation.ARTIFACT_TEXT: 128 | ext = ".pb.json" 129 | contents = MessageToJson(artifact).encode("utf-8") 130 | else: 131 | ext = ".pb" 132 | contents = artifact.SerializeToString() 133 | out_p = f"{artifact_p}{ext}" 134 | if write: 135 | with open(out_p, "wb") as f: 136 | f.write(bytes(contents)) 137 | if verbose: 138 | artifact_t = generation.ArtifactType.Name(artifact.type) 139 | logger.info(f"wrote {artifact_t} to {out_p}") 140 | 141 | yield [out_p, artifact] 142 | idx += 1 143 | 144 | 145 | def open_images( 146 | images: Union[ 147 | Sequence[Tuple[str, generation.Artifact]], 148 | Generator[Tuple[str, generation.Artifact], None, None], 149 | ], 150 | verbose: bool = False, 151 | ) -> Generator[Tuple[str, generation.Artifact], None, None]: 152 | """ 153 | Open the images from the filenames and Artifacts tuples. 154 | 155 | :param images: The tuples of Artifacts and associated images to open. 156 | :return: A Generator of tuples of image filenames and Artifacts, intended 157 | for passthrough. 158 | """ 159 | from PIL import Image 160 | 161 | for path, artifact in images: 162 | if artifact.type == generation.ARTIFACT_IMAGE: 163 | if verbose: 164 | logger.info(f"opening {path}") 165 | img = Image.open(io.BytesIO(artifact.binary)) 166 | img.show() 167 | yield [path, artifact] 168 | 169 | 170 | class StabilityInference: 171 | def __init__( 172 | self, 173 | host: str = "grpc.stability.ai:443", 174 | key: str = f"{STABILITY_API_KEY}", 175 | engine: str = "stable-diffusion-v1", 176 | verbose: bool = False, 177 | wait_for_ready: bool = True, 178 | ): 179 | """ 180 | Initialize the client. 181 | 182 | :param host: Host to connect to. 183 | :param key: Key to use for authentication. 184 | :param engine: Engine to use. 185 | :param verbose: Whether to print debug messages. 186 | :param wait_for_ready: Whether to wait for the server to be ready, or 187 | to fail immediately. 188 | """ 189 | self.verbose = verbose 190 | self.engine = engine 191 | 192 | self.grpc_args = {"wait_for_ready": wait_for_ready} 193 | 194 | if verbose: 195 | logger.info(f"Opening channel to {host}") 196 | 197 | call_credentials = [] 198 | 199 | if host.endswith("443"): 200 | if key: 201 | call_credentials.append( 202 | grpc.access_token_call_credentials(f"{key}")) 203 | else: 204 | raise ValueError(f"key is required for {host}") 205 | channel_credentials = grpc.composite_channel_credentials( 206 | grpc.ssl_channel_credentials(), *call_credentials 207 | ) 208 | channel = grpc.secure_channel(host, channel_credentials) 209 | else: 210 | if key: 211 | logger.warning( 212 | "Not using authentication token due to non-secure transport" 213 | ) 214 | channel = grpc.insecure_channel(host) 215 | 216 | if verbose: 217 | logger.info(f"Channel opened to {host}") 218 | self.stub = generation_grpc.GenerationServiceStub(channel) 219 | 220 | def generate( 221 | self, 222 | prompt: Union[List[str], str], 223 | height: int = 512, 224 | width: int = 512, 225 | cfg_scale: float = 7.0, 226 | sampler: generation.DiffusionSampler = generation.SAMPLER_K_LMS, 227 | steps: int = 50, 228 | seed: Union[Sequence[int], int] = 0, 229 | samples: int = 1, 230 | safety: bool = True, 231 | classifiers: generation.ClassifierParameters = None, 232 | ) -> Generator[generation.Answer, None, None]: 233 | """ 234 | Generate images from a prompt. 235 | 236 | :param prompt: Prompt to generate images from. 237 | :param height: Height of the generated images. 238 | :param width: Width of the generated images. 239 | :param cfg_scale: Scale of the configuration. 240 | :param sampler: Sampler to use. 241 | :param steps: Number of steps to take. 242 | :param seed: Seed for the random number generator. 243 | :param samples: Number of samples to generate. 244 | :param safety: Whether to use safety mode. 245 | :param classifications: Classifier parameters to use. 246 | :return: Generator of Answer objects. 247 | """ 248 | if safety and classifiers is None: 249 | classifiers = generation.ClassifierParameters() 250 | 251 | if not prompt: 252 | raise ValueError("prompt must be provided") 253 | 254 | request_id = str(uuid.uuid4()) 255 | 256 | if not seed: 257 | seed = [random.randrange(0, 4294967295)] 258 | 259 | if isinstance(prompt, str): 260 | prompt = [generation.Prompt(text=prompt)] 261 | else: 262 | prompt = [generation.Prompt(text=p) for p in prompt] 263 | 264 | rq = generation.Request( 265 | engine_id=self.engine, 266 | request_id=request_id, 267 | prompt=prompt, 268 | image=generation.ImageParameters( 269 | transform=generation.TransformType(diffusion=sampler), 270 | height=height, 271 | width=width, 272 | seed=seed, 273 | steps=steps, 274 | samples=samples, 275 | parameters=[ 276 | generation.StepParameter( 277 | scaled_step=0, 278 | sampler=generation.SamplerParameters(cfg_scale=cfg_scale), 279 | ) 280 | ], 281 | ), 282 | classifier=classifiers, 283 | ) 284 | 285 | if self.verbose: 286 | logger.info("Sending request.") 287 | 288 | start = time.time() 289 | for answer in self.stub.Generate(rq, **self.grpc_args): 290 | duration = time.time() - start 291 | if self.verbose: 292 | if len(answer.artifacts) > 0: 293 | artifact_ts = [ 294 | generation.ArtifactType.Name(artifact.type) 295 | for artifact in answer.artifacts 296 | ] 297 | logger.info( 298 | f"Got {answer.answer_id} with {artifact_ts} in " 299 | f"{duration:0.2f}s" 300 | ) 301 | else: 302 | logger.info( 303 | f"Got keepalive {answer.answer_id} in " 304 | f"{duration:0.2f}s" 305 | ) 306 | 307 | yield answer 308 | start = time.time() 309 | 310 | 311 | def build_request_dict(cli_args: Namespace) -> Dict[str, Any]: 312 | """ 313 | Build a Request arguments dictionary from the CLI arguments. 314 | """ 315 | return { 316 | "height": cli_args.height, 317 | "width": cli_args.width, 318 | "cfg_scale": cli_args.cfg_scale, 319 | "sampler": get_sampler_from_str(cli_args.sampler), 320 | "steps": cli_args.steps, 321 | "seed": cli_args.seed, 322 | "samples": cli_args.num_samples, 323 | } 324 | 325 | #import bpy 326 | import bpy 327 | import numpy as np 328 | 329 | from PIL import Image 330 | from stability_sdk import client 331 | import stability_sdk.interfaces.gooseai.generation.generation_pb2 as generation 332 | import base64 333 | import zlib 334 | import struct 335 | 336 | #You will also need to put your API key here 337 | stability_api = client.StabilityInference( 338 | key=f"{STABILITY_API_KEY}", 339 | verbose=True, 340 | ) 341 | 342 | 343 | 344 | def pil_to_image(pil_image, name='NewImage'): 345 | ''' 346 | PIL image pixels is 2D array of byte tuple (when mode is 'RGB', 'RGBA') or byte (when mode is 'L') 347 | bpy image pixels is flat array of normalized values in RGBA order 348 | ''' 349 | # setup PIL image reading 350 | width = pil_image.width 351 | height = pil_image.height 352 | pil_pixels = pil_image.load() 353 | byte_to_normalized = 1.0 / 255.0 354 | num_pixels = width * height 355 | # setup bpy image 356 | channels = 4 357 | bpy_image = bpy.data.images.new(name, width=width, height=height) 358 | # bpy image has a flat RGBA array (similar to JS Canvas) 359 | bpy_image.pixels = (np.asarray(pil_image.convert('RGBA'),dtype=np.float32) * byte_to_normalized).ravel() 360 | return bpy_image 361 | 362 | # Save the image to the current working directory 363 | def save_image(image, name): 364 | img = pil_to_image(image) 365 | image.save(f"{name}.png") 366 | 367 | def parse_response(response): 368 | for resp in response: 369 | for artifact in resp.artifacts: 370 | print('artifact type: '+ str(artifact.type)) 371 | if artifact.type == 1: 372 | img = Image.open(io.BytesIO(artifact.binary)) 373 | return img 374 | 375 | 376 | # Create a new image from a prompt and save it to the current working directory 377 | # if no name is provided, save the image the text of the prompt 378 | def generate_texture_from_prompt(prompt,size,name): 379 | # if name is not defined, use the prompt as name 380 | if name is None: 381 | name = prompt 382 | # send request to stability api 383 | request = stability_api.generate(prompt, width=size[0], height=size[1]) 384 | # parse response (PIL image is returned) 385 | image = parse_response(request) 386 | # save image to current working directory 387 | save_image(image, name) 388 | 389 | 390 | return request 391 | 392 | generate_texture_from_prompt(prompt="A painting of luminescent translucent cast glass 3d xray abstract supernova portrait made of many vividly colored neon tube tribal masks filled with glowing uv blacklight led light trails and curved lightsaber glowing graffiti by okuda san miguel and kandinsky in a cubist face, in metallic light trail calligraphic light saber galaxies by okuda san miguel and kandinsky on a starry black canvas, galaxy gas brushstrokes, metallic flecked paint, metallic flecks, glittering metal paint, metallic paint, glossy flecks of iridescence, glow in the dark, uv, blacklight, uv blacklight, 8k, 4k, brush strokes, painting, highly detailed, iridescent texture, brushed metal", size=(512,512), name="DiffusionTexture_001") 393 | -------------------------------------------------------------------------------- /StableDiffusionWebUI.py: -------------------------------------------------------------------------------- 1 | # Stablender Diffusion 0.0.1 2 | # author: @shellworld 3 | # license: MIT 4 | 5 | # imports 6 | import uuid 7 | import bpy 8 | import os 9 | import sys 10 | import requests 11 | import json 12 | from base64 import b64decode, b64encode 13 | import io 14 | from PIL import Image 15 | import cv2 16 | import numpy as np 17 | from datetime import datetime 18 | from enum import Enum 19 | 20 | # constants 21 | # function indices 22 | class FunctionIndex(Enum): 23 | TEXT_TO_IMAGE = 4 24 | IMAGE_TO_IMAGE = 17 25 | IMAGE_TO_IMAGE_WITH_MASK = 16 26 | 27 | # set vars 28 | # comment out in operators for accepting input variables 29 | 30 | # replace with gradio url 31 | url = 'https://#####.gradio.app' 32 | prompt = "landscape painting, oil on cavas, high detail 4k" 33 | cfg_scale = 7.5 34 | width = 512 35 | height = 512 36 | steps = 50 37 | num_batches = 1 38 | batch_size = 1 39 | strength = .71 40 | 41 | # set headers for requests to StableDiffusion Web Ui 42 | headers = { 43 | "authority": f"{url}", 44 | "method": "POST", 45 | "path": "/api/predict/", 46 | "scheme": "https", 47 | "accept": "*/*", 48 | "accept-encoding": "gzip, deflate, br", 49 | "accept-language": "en-US,en;q=0.9", 50 | "dnt": "1", 51 | "origin": f"{url}", 52 | "referer": f"{url}", 53 | "sec-ch-ua": "`\"Chromium`\";v=\"104\", `\" Not A;Brand`\";v=\"99\", `\"Google Chrome`\";v=\"104\"\"", 54 | "sec-ch-ua-mobile": "?0", 55 | "sec-ch-ua-platform": "`\"Windows`\"", 56 | "sec-fetch-dest": "empty", 57 | "sec-fetch-mode": "cors", 58 | "sec-fetch-site": "same-origin" 59 | } 60 | 61 | # constant static variables 62 | MAX_STEPS = 400 63 | MIN_STEPS = 10 64 | MAX_PROMPT_LENGTH = 1000 65 | CFG_SCALE_FLOOR = 0.1 66 | CFG_SCALE_CEILING = 10.0 67 | SIZE_MIN = 512 68 | SIZE_MAX = 768 69 | HIGH_BATCH_NUMBER = 16 70 | HIGH_BATCH_NUMBER_ENABLED = False 71 | MAX_LENGTH_FILE_NAME=24 72 | 73 | 74 | # predict 75 | def predict(prompt:str, steps:int, cfg_scale:float, width:int, height:int, num_batches=1, batch_size=1, image_string="", mask_string="", function_index=FunctionIndex.TEXT_TO_IMAGE,strength=.50): 76 | # get prediction 77 | 78 | # set function index based on input 79 | # this will grab the corresponding enum value 80 | fn_index = function_index.value 81 | 82 | 83 | # define validation functions 84 | 85 | def checkPrompt(prompt): 86 | #prompt is not null or empty 87 | if prompt is None or prompt == "": 88 | raise (Exception("prompt is required")) 89 | # prompt is a valid string 90 | if type(prompt) != str: 91 | raise (Exception("prompt must be a string")) 92 | # prompt is not too long 93 | if len(prompt) > MAX_PROMPT_LENGTH: 94 | raise ( 95 | Exception(f"prompt must be less than {str(MAX_PROMPT_LENGTH)} characters")) 96 | 97 | def checkSteps(steps): 98 | if steps is None: 99 | raise (Exception("steps is required")) 100 | if type(steps) != int: 101 | raise (Exception("steps must be an integer")) 102 | if steps < MIN_STEPS: 103 | raise (Exception(f"steps must be greater than {str(MIN_STEPS)}")) 104 | if steps > MAX_STEPS: 105 | #steps is not null or empty 106 | if steps is None or steps == "": 107 | raise (Exception("steps are required")) 108 | # steps is a valid integer 109 | if type(steps) != int: 110 | raise (Exception("steps must be an integer")) 111 | # steps are within the range of 10 to 400 112 | if steps < MIN_STEPS or steps > MAX_STEPS: 113 | raise ( 114 | Exception(f"steps must be between {str(MIN_STEPS)} and {str(MAX_STEPS)}")) 115 | 116 | def checkCfgScale(cfg_scale): 117 | #cfg_scale is not null or empty 118 | if cfg_scale is None or cfg_scale == "": 119 | raise (Exception("cfg_scale is required")) 120 | # cfg_scale is a valid float 121 | if type(cfg_scale) != float: 122 | raise (Exception("cfg_scale must be a float")) 123 | # cfg_scale is within the range of -CFG_SCALE_FLOOR to CFG_SCALE_CEILING 124 | if cfg_scale < CFG_SCALE_FLOOR or cfg_scale > CFG_SCALE_CEILING: 125 | raise (Exception( 126 | f"cfg_scale must be between {str(CFG_SCALE_FLOOR)} and {str(CFG_SCALE_CEILING)}")) 127 | 128 | def checkWidth(width): 129 | #width is not null or empty 130 | if width is None or width == "": 131 | raise (Exception("width is required")) 132 | # width is a valid integer 133 | if type(width) != int: 134 | raise (Exception("width must be an integer")) 135 | # width is greater than SIZE_MIN and less than SIZE_MAX 136 | if width < SIZE_MIN or width > SIZE_MAX: 137 | raise ( 138 | Exception(f"width must be between {str(SIZE_MIN)}px and {str(SIZE_MAX)}px")) 139 | 140 | def checkHeight(height): 141 | #height is not null or empty 142 | if height is None or height == "": 143 | raise (Exception("height is required")) 144 | # height is a valid integer 145 | if type(height) != int: 146 | raise (Exception("height must be an integer")) 147 | # height is greater than SIZE_MIN and less than SIZE_MAX 148 | if height < SIZE_MIN or height > SIZE_MAX: 149 | raise ( 150 | Exception(f"height must be between {str(SIZE_MIN)}px and {str(SIZE_MAX)}px")) 151 | 152 | def checkNumBatches(num_batches): 153 | #num_batches is not null or empty 154 | if num_batches is None or num_batches == "": 155 | raise (Exception("num_batches is required")) 156 | # num_batches is a valid integer greater than 0 157 | if type(num_batches) != int or num_batches <= 0: 158 | raise (Exception("num_batches must be an integer greater than zero")) 159 | # num_batches is greater than 16 160 | if num_batches > HIGH_BATCH_NUMBER and not HIGH_BATCH_NUMBER_ENABLED: 161 | raise (Exception( 162 | f"num_batches must be less than {str(HIGH_BATCH_NUMBER)} or set HIGH_BATCH_NUMBER_ENABLED to True")) 163 | 164 | # Developer Note: 165 | # It might make sense to disable the batch size option 166 | # in your interface. 167 | 168 | def checkBatchSize(batch_size): 169 | if batch_size is None or batch_size == "": 170 | raise (Exception("batch_size is required")) 171 | # batch_size is a valid integer greater than 0 172 | if type(batch_size) != int or batch_size <= 0: 173 | raise (Exception("batch_size must be an integer greater than zero")) 174 | # batch_size is greater than 1 175 | if batch_size > 1: 176 | raise (Exception(f"batch_size must be set to one for now.")) 177 | 178 | def checkImageString(image_string): 179 | if image_string is not None or image_string != "": 180 | # image_string is a valid string 181 | if type(image_string) != str: 182 | print(type(image_string)) 183 | raise (Exception(f"image_string must be a string, but it is currently {type(image_string)}")) 184 | # image string is a valid base64 png datastream 185 | if not image_string.startswith("data:image/png;base64,"): 186 | print(f"DEBUG:{image_string}") 187 | raise (Exception("image_string must be a valid base64 png datastream")) 188 | 189 | 190 | def checkMaskString(mask_string): 191 | if mask_string is not None or mask_string != "": 192 | # mask_string is a valid string 193 | if type(mask_string) != str: 194 | print(type(image_string)) 195 | raise (Exception("mask_string must be a string")) 196 | # mask string is a valid base64 png datastream 197 | if not mask_string.startswith("data:image/png;base64,"): 198 | print(f"DEBUG:{image_string}") 199 | raise (Exception("mask_string must be a valid base64 png datastream")) 200 | 201 | def checkStrength(strength): 202 | if strength is not None or strength != "": 203 | # strength is a valid float 204 | if type(strength) != float: 205 | raise (Exception("strength must be a float")) 206 | # strength is within the range of 0 to 1 207 | if strength < 0 or strength > 1: 208 | raise (Exception(f"strength must be between 0 and 1")) 209 | 210 | def validate(prompt, steps, cfg_scale, width, height, num_batches, batch_size, image_string, mask_string, strength): 211 | checkPrompt(prompt) 212 | checkSteps(steps) 213 | checkCfgScale(cfg_scale) 214 | checkWidth(width) 215 | checkHeight(height) 216 | checkNumBatches(num_batches) 217 | checkBatchSize(batch_size) 218 | if function_index == FunctionIndex.IMAGE_TO_IMAGE or function_index == FunctionIndex.IMAGE_TO_IMAGE_WITH_MASK: 219 | checkImageString(image_string) 220 | checkStrength(strength) 221 | if function_index== FunctionIndex.IMAGE_TO_IMAGE_WITH_MASK: 222 | checkMaskString(mask_string) 223 | 224 | def validateResponse(response): 225 | if response is None or response == "": 226 | raise (Exception(f"Something went wrong: {response}")) 227 | if response.status_code != 200: 228 | raise (Exception(f"Something went wrong: {response}")) 229 | 230 | # validate the input 231 | validate(prompt, steps, cfg_scale, width, height, num_batches, 232 | batch_size, image_string, mask_string, strength) 233 | 234 | # depending on which mode we are in, we need to set the data differently 235 | 236 | # txt2img 237 | if function_index == FunctionIndex.TEXT_TO_IMAGE: 238 | data = [prompt, steps, "k_lms", 239 | ["Normalize Prompt Weights (ensure sum of weights add up to 1.0)", 240 | "Save individual images", 241 | "Save grid", 242 | "Sort samples by prompt", 243 | "Write sample info files"], 244 | "RealESRGAN_x4plus", 0, num_batches, batch_size, 245 | cfg_scale, "", width, height, None, 0, "" 246 | ] 247 | 248 | # img2img 249 | elif function_index == FunctionIndex.IMAGE_TO_IMAGE: 250 | data = [prompt, "Crop", image_string, "Keep masked area", 251 | 3, steps, "k_lms", 252 | ["Normalize Prompt Weights (ensure sum of weights add up to 1.0)", 253 | "Save individual images", 254 | "Save grid", 255 | "Sort samples by prompt", 256 | "Write sample info files"], 257 | "RealESRGAN_x4plus", 258 | num_batches, batch_size, cfg_scale, strength, 259 | None, width, height, "Just resize", None 260 | ] 261 | 262 | # img2img with mask 263 | elif function_index == FunctionIndex.IMAGE_TO_IMAGE_WITH_MASK: 264 | data = [prompt, "Mask", {"image": image_string, "mask": mask_string}, 265 | "Keep masked area", 3, steps, "k_lms", 266 | ["Normalize Prompt Weights (ensure sum of weights add up to 1.0)", 267 | "Save individual images", 268 | "Save grid", 269 | "Sort samples by prompt", 270 | "Write sample info files"], 271 | "RealESRGAN_x4plus", 272 | num_batches, batch_size, cfg_scale, strength, 273 | None, width, height, "Just resize", None 274 | ] 275 | 276 | response = requests.post(url + '/api/predict/', headers=headers, json={ 277 | "fn_index": fn_index, 278 | "data":data 279 | }) 280 | 281 | # validate the response 282 | validateResponse(response) 283 | 284 | return response.json() 285 | 286 | 287 | 288 | 289 | # utility classes 290 | 291 | # function that converts png datastream to Image 292 | def stringToRGB(base64_string:str): 293 | header, encoded = base64_string.split(",", 1) 294 | imgdata = b64decode(encoded) 295 | img = Image.open(io.BytesIO(imgdata)) 296 | return img 297 | 298 | # function that converts Image to png datastream 299 | def rgbToString(image:Image.Image): 300 | img = image.pixels 301 | img = Image.fromarray(img.reshape(image.size[1], image.size[0], 4)) 302 | img_bytes = io.BytesIO() 303 | img.save(img_bytes, format='PNG') 304 | img_bytes.seek(0) 305 | img_string = b64encode(img_bytes.read()).decode('ascii') 306 | return img_string 307 | 308 | 309 | def parseResults(response): 310 | #return array of b64 strings from response[data][0] 311 | 312 | # get the data property's zeroth index 313 | data = response["data"][0] 314 | 315 | #if the length of data is zero, raise an exception 316 | if len(data) == 0: 317 | raise (Exception("No data returned")) 318 | 319 | elif len(data) == 1: 320 | # if the length of data is one, then we have a single image 321 | # return the image 322 | return [data[0]] 323 | 324 | else: 325 | images = [] 326 | for i in range(len(data)): 327 | # if data>1, the first image will be a grid, ignore this image 328 | if i == 0: 329 | continue 330 | # return multiple images 331 | else: 332 | images.append(data[i]) 333 | return images 334 | 335 | def convertPNGDatastreamsToBPYImages(base64_strings:str): 336 | images = [] 337 | #convert to CV2 images 338 | for base64_string in base64_strings: 339 | images.append(stringToRGB(base64_string)) 340 | #now convert the cv2 images to bpy images 341 | for i in range(len(images)): 342 | 343 | images[i] = bpy.data.images.new(f"image_{i}", images[i].size[0], images[i].size[1], alpha=True) 344 | images[i].pixels = images[i].pixels[:] 345 | images[i].filepath_raw = f"image_{i}.png" 346 | images[i].file_format = 'PNG' 347 | return images 348 | 349 | def convertBPYImageToBase64PNGDataStream(referenceToBPYImage): 350 | selectedImage = bpy.data.images[referenceToBPYImage] 351 | 352 | #convert the bpy image to a cv2 image 353 | img = cv2.imread(f"{bpy.path.abspath(selectedImage.filepath_raw)}") 354 | #convert the cv2 image to a b64 string 355 | _, im_arr = cv2.imencode('.jpg', img) # im_arr: image in Numpy one-dim array format. 356 | im_bytes = im_arr.tobytes() 357 | im_b64 = b64encode(im_bytes).decode('ascii') 358 | im_b64 = "data:image/png;base64," + im_b64 359 | return im_b64 360 | 361 | 362 | def pil_to_image(pil_image, name='NewImage'): 363 | ''' 364 | PIL image pixels is 2D array of byte tuple (when mode is 'RGB', 'RGBA') or byte (when mode is 'L') 365 | bpy image pixels is flat array of normalized values in RGBA order 366 | ''' 367 | # setup PIL image reading 368 | 369 | # swap red and blue channels 370 | pil_image = pil_image.transpose(Image.FLIP_TOP_BOTTOM) 371 | # convert to bpy image 372 | 373 | 374 | width = pil_image.width 375 | height = pil_image.height 376 | pil_pixels = pil_image.load() 377 | byte_to_normalized = 1.0 / 255.0 378 | num_pixels = width * height 379 | # setup bpy image 380 | channels = 4 381 | bpy_image = bpy.data.images.new(name, width=width, height=height) 382 | 383 | 384 | # bpy image has a flat RGBA array (similar to JS Canvas) 385 | bpy_image.pixels = (np.asarray(pil_image.convert('RGBA'),dtype=np.float32) * byte_to_normalized).ravel() 386 | return bpy_image 387 | 388 | #save bpy images 389 | def saveImage(image_data, filename): 390 | img = stringToRGB(image_data) 391 | img.save(filename) 392 | return img 393 | 394 | # generate a safe filename 395 | def generateSafeNameFromPromptAndIndex(prompt, index): 396 | prompt = prompt.replace(" ", "_") 397 | prompt = prompt.replace(",", "-") 398 | prompt = prompt[0:MAX_LENGTH_FILE_NAME] 399 | return prompt + "_" + str(index) 400 | 401 | ## MAIN FUNCTIONS 402 | 403 | # function that converts base64 png datastream to Image 404 | def requestImg(prompt, steps, cfg_scale, width, height, fn, bpy_image=None, mask_image=None, batch_num=1, batch_size=1, strength=50): 405 | #request the image 406 | response = predict(prompt, steps, cfg_scale, width, height, 407 | num_batches, batch_size,bpy_image, mask_image, fn) 408 | 409 | #results will be a list of base64 strings 410 | results = parseResults(response) 411 | img = saveImage(results[0], f"{generateSafeNameFromPromptAndIndex(prompt, 0)}.png") 412 | #add image to bpy.data.images 413 | blender_image = pil_to_image(img, f"{generateSafeNameFromPromptAndIndex(prompt, 0)}") 414 | return f"{generateSafeNameFromPromptAndIndex(prompt, 0)}" 415 | 416 | 417 | #Text to image example 418 | requestImg(prompt, steps, cfg_scale, width, height, FunctionIndex.TEXT_TO_IMAGE) 419 | 420 | #Image to image example (replace V with the name of the blender texture) 421 | #requestImg(prompt, steps, cfg_scale, width, height, FunctionIndex.IMAGE_TO_IMAGE, bpy_image=convertBPYImageToBase64PNGDataStream("V"), strength=strength) 422 | 423 | --------------------------------------------------------------------------------