Upload stable_diffusion_loader.py

stable_diffusion_loader.py

@@ -0,0 +1,130 @@
+import os
+import torch
+import base64
+from io import BytesIO
+from PIL import Image
+
+from transformers import CLIPTokenizer, CLIPTextModel
+from diffusers import (
+    StableDiffusionPipeline,
+    UNet2DConditionModel,
+    AutoencoderKL,
+    DDIMScheduler,
+)
+
+
+def load_custom_pipeline(
+    model_path: str = "./fine-tuned-model",
+    use_mps_if_available: bool = True
+):
+    """
+    Loads your custom fine-tuned Stable Diffusion model from a local folder structure.
+    Returns a pipeline object ready for inference.
+    """
+    # Load tokenizer
+    tokenizer = CLIPTokenizer.from_pretrained(os.path.join(model_path, "tokenizer"))
+
+    # Load text encoder
+    text_encoder = CLIPTextModel.from_pretrained(
+        os.path.join(model_path, "text_encoder"),
+        torch_dtype=torch.float32
+    )
+
+    # Load UNet
+    unet = UNet2DConditionModel.from_pretrained(
+        os.path.join(model_path, "unet"),
+        torch_dtype=torch.float32
+    )
+
+    # Load VAE
+    vae = AutoencoderKL.from_pretrained(
+        os.path.join(model_path, "vae"),
+        torch_dtype=torch.float32
+    )
+
+    # Load scheduler
+    scheduler = DDIMScheduler.from_pretrained(
+        "CompVis/stable-diffusion-v1-4",
+        subfolder="scheduler"
+    )
+
+    # Create the pipeline
+    pipe = StableDiffusionPipeline(
+        tokenizer=tokenizer,
+        text_encoder=text_encoder,
+        vae=vae,
+        unet=unet,
+        scheduler=scheduler,
+        safety_checker=None,   # Disable safety checker
+        feature_extractor=None
+    )
+
+    # Set device
+    device = (
+        torch.device("mps") if (torch.backends.mps.is_available() and use_mps_if_available)
+        else torch.device("cpu")
+    )
+    pipe.to(device)
+
+    # Optional: reduce memory usage
+    pipe.enable_attention_slicing()
+
+    return pipe
+
+
+def load_base_pipeline(
+    model_id: str = "CompVis/stable-diffusion-v1-4",
+    use_mps_if_available: bool = True
+):
+    """
+    Loads the original Stable Diffusion v1.4 model from Hugging Face.
+    Returns a pipeline object ready for inference.
+    """
+    pipe = StableDiffusionPipeline.from_pretrained(
+        model_id,
+        torch_dtype=torch.float32,
+        safety_checker=None,
+        feature_extractor=None
+    )
+    device = (
+        torch.device("mps") if (torch.backends.mps.is_available() and use_mps_if_available)
+        else torch.device("cpu")
+    )
+    pipe.to(device)
+    pipe.enable_attention_slicing()
+    return pipe
+
+
+def generate_image(
+    pipe: StableDiffusionPipeline,
+    prompt: str,
+    num_inference_steps: int = 50,
+    guidance_scale: float = 7.5,
+    seed: int = None
+):
+    """
+    Generates a single image from the provided pipeline and prompt.
+    Optionally accepts a 'seed' for reproducibility.
+    """
+    if seed is not None:
+        generator = torch.Generator(device=pipe.device).manual_seed(seed)
+    else:
+        generator = None
+
+    with torch.no_grad():
+        result = pipe(
+            prompt=prompt,
+            num_inference_steps=num_inference_steps,
+            guidance_scale=guidance_scale,
+            generator=generator
+        )
+    return result.images[0]
+
+
+def pil_image_to_base64_str(img: Image.Image) -> str:
+    """
+    Converts a PIL Image into a Base64-encoded PNG string.
+    """
+    buffered = BytesIO()
+    img.save(buffered, format="PNG")
+    return base64.b64encode(buffered.getvalue()).decode("utf-8")