Spaces:
Sleeping
Sleeping
| import os | |
| import io | |
| import warnings | |
| from PIL import Image | |
| from stability_sdk import client | |
| import stability_sdk.interfaces.gooseai.generation.generation_pb2 as generation | |
| import uuid | |
| import gradio as gr | |
| # Our Host URL should not be prepended with "https" nor should it have a trailing slash. | |
| os.environ["STABILITY_HOST"] = "grpc.stability.ai:443" | |
| def get_image(prompt, api_key_stability_ai): | |
| # Sign up for an account at the following link to get an API Key. | |
| # https://platform.stability.ai/ | |
| # Click on the following link once you have created an account to be taken to your API Key. | |
| # https://platform.stability.ai/account/keys | |
| # Set up our connection to the API. | |
| if api_key_stability_ai == "": | |
| raise gr.Error("Please add your Stability AI API key ") | |
| else: | |
| try: | |
| stability_api = client.StabilityInference( | |
| key=api_key_stability_ai, # API Key reference. | |
| verbose=True, # Print debug messages. | |
| engine="stable-diffusion-xl-1024-v1-0", # Set the engine to use for generation. | |
| # Check out the following link for a list of available engines: https://platform.stability.ai/docs/features/api-parameters#engine | |
| ) | |
| # Set up our initial generation parameters. | |
| answers = stability_api.generate( | |
| prompt=prompt, # The prompt we want to generate an image from. | |
| seed=4253978046, # If a seed is provided, the resulting generated image will be deterministic. | |
| # What this means is that as long as all generation parameters remain the same, you can always recall the same image simply by generating it again. | |
| # Note: This isn't quite the case for Clip Guided generations, which we'll tackle in a future example notebook. | |
| steps=30, # Amount of inference steps performed on image generation. Defaults to 30. | |
| cfg_scale=8.0, # Influences how strongly your generation is guided to match your prompt. | |
| # Setting this value higher increases the strength in which it tries to match your prompt. | |
| # Defaults to 7.0 if not specified. | |
| width=512, # Generation width, defaults to 512 if not included. | |
| height=512, # Generation height, defaults to 512 if not included. | |
| samples=1, # Number of images to generate, defaults to 1 if not included. | |
| sampler=generation.SAMPLER_K_DPMPP_2M, # Choose which sampler we want to denoise our generation with. | |
| # Defaults to k_dpmpp_2m if not specified. Clip Guidance only supports ancestral samplers. | |
| # (Available Samplers: ddim, plms, k_euler, k_euler_ancestral, k_heun, k_dpm_2, k_dpm_2_ancestral, k_dpmpp_2s_ancestral, k_lms, k_dpmpp_2m, k_dpmpp_sde) | |
| ) | |
| # print("Finish the prompt") | |
| # Set up our warning to print to the console if the adult content classifier is tripped. | |
| # If adult content classifier is not tripped, save generated images. | |
| for resp in answers: | |
| for artifact in resp.artifacts: | |
| # if artifact.finish_reason == generation.FILTER: | |
| # print(artifact.finish_reason) | |
| # print("Warning") | |
| # warnings.warn( | |
| # "Your request activated the API's safety filters and could not be processed." | |
| # "Please modify the prompt and try again.") | |
| if artifact.type == generation.ARTIFACT_IMAGE: | |
| img = Image.open(io.BytesIO(artifact.binary)) | |
| unique_filename = str(uuid.uuid4()) | |
| img.save( | |
| str(unique_filename) + ".png" | |
| ) # Save our generated images with their seed number as the filename. | |
| return unique_filename + ".png" | |
| except Exception as error: | |
| print(str(error)) | |
| raise gr.Error( | |
| "An error occurred while generating the image. Please try again." | |
| ) | |