WIP- Under maintenence- Web research working.

lib/text_to_image/stabl_diff_img2html.py

@@ -0,0 +1,66 @@
+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
+
+# Set the host URL environment variable. Ensure it doesn't have 'https' or a trailing slash.
+os.environ['STABILITY_HOST'] = 'grpc.stability.ai:443'
+
+# Ensure you sign up for an account to obtain an API key:
+# https://platform.stability.ai/
+# Your API key can be found here after account creation:
+# https://platform.stability.ai/account/keys
+
+
+def generate_stable_diffusion_image(prompt, image_dir):
+    """
+    Generate images using Stable Diffusion API based on a given prompt.
+
+    Args:
+        prompt (str): The prompt to generate the image.
+        image_dir (str): The directory where the image will be saved.
+
+    Raises:
+        Warning: If the adult content classifier is triggered.
+        Exception: For any issues during image generation or saving.
+    """
+    try:
+        # Initialize the StabilityInference client with the API key and other settings.
+        stability_api = client.StabilityInference(
+            key=os.environ['STABILITY_KEY'],  # Reference to the API key.
+            verbose=True,  # Enable verbose mode for debug messages.
+            engine="stable-diffusion-xl-1024-v1-0",  # Engine used for generation.
+        )
+
+        # Generating the image with specified parameters.
+        answers = stability_api.generate(
+            prompt=prompt,
+            seed=4253978046,  # Deterministic seed for reproducible results.
+            steps=50,  # Number of inference steps.
+            cfg_scale=7.0,  # Strength of prompt matching.
+            width=1024, height=1024,  # Image dimensions.
+            samples=1,  # Number of images to generate.
+            sampler=generation.SAMPLER_K_DPMPP_2M  # Denoising sampler selection.
+        )
+
+        # Process responses and save images.
+        for resp in answers:
+            for artifact in resp.artifacts:
+                if artifact.finish_reason == generation.FILTER:
+                    warnings.warn(
+                        "Request activated safety filters. Modify the prompt and retry."
+                    )
+                if artifact.type == generation.ARTIFACT_IMAGE:
+                    img = Image.open(io.BytesIO(artifact.binary))
+                    img_name = os.path.join(image_dir, f"{artifact.seed}.png")
+                    img.show()
+                    img.save(img_name)  # Save the image with the seed in the filename.
+
+    except Exception as e:
+        raise Exception(f"Error during image generation or saving: {e}")
+
+# Example usage:
+# generate_stable_diffusion_image("A futuristic cityscape", "/path/to/save/images/")
+