glowforge-dev/stable-diffusion-2-1-img2img-jonathan2

1from typing import  Dict, List, Any2import torch3from PIL import Image4from io import BytesIO5from diffusers import StableDiffusionPipeline, StableDiffusionImg2ImgPipeline, DDIMScheduler6from transformers.utils import logging7 8import base649import requests10from io import BytesIO11from PIL import Image12 13logging.set_verbosity_info()14logger = logging.get_logger("transformers")15 16def load_image(image_url):17    if image_url.startswith('data:'):18        # Decode base64 data_uri19        image_data = base64.b64decode(image_url.split(',')[1])20        image = Image.open(BytesIO(image_data))21    else:22        # Load standard image url23        response = requests.get(image_url)24        image = Image.open(BytesIO(response.content))25    return image26 27# set device28device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')29 30if device.type != 'cuda':31    raise ValueError("need to run on GPU")32 33model_id = "stabilityai/stable-diffusion-2-1-base"34 35class EndpointHandler():36    def __init__(self, path=""):37        # load the optimized model38        self.textPipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)39        self.textPipe.scheduler = DDIMScheduler.from_config(self.textPipe.scheduler.config)40        self.textPipe = self.textPipe.to(device)41 42        # create an img2img model43        self.imgPipe = StableDiffusionImg2ImgPipeline.from_pretrained(model_id, torch_dtype=torch.float16)44        self.imgPipe.scheduler = DDIMScheduler.from_config(self.imgPipe.scheduler.config)45        self.imgPipe = self.imgPipe.to(device)46 47    def __call__(self, data: Any) -> List[List[Dict[str, float]]]:48        """49        Args:50            data (:obj:):51                includes the input data and the parameters for the inference.52        Return:53            A :obj:`dict`:. base64 encoded image54        """55        prompt = data.pop("inputs", data)56        url = data.pop("url", data)57 58        init_image = load_image(url).convert("RGB")59        init_image.thumbnail((512, 512))60 61 62        params = data.pop("parameters", data)63 64        # hyperparamters65        num_inference_steps = params.pop("num_inference_steps", 25)66        guidance_scale = params.pop("guidance_scale", 7.5)67        negative_prompt = params.pop("negative_prompt", None)68        height = params.pop("height", None)69        strength = params.pop("strength", 0.8)70        width = params.pop("width", None)71        manual_seed = params.pop("manual_seed", -1)72        logger.info(f"strength: {strength}, manual_seed: {manual_seed}, inference_steps: {num_inference_steps}, guidance_scale: {guidance_scale}")73        out = None74 75        generator = torch.Generator(device='cuda')76        generator.manual_seed(manual_seed)77        # run img2img pipeline78        out = self.imgPipe(prompt,79                    image=init_image,80                    strength=strength,81                    num_inference_steps=num_inference_steps,82                    guidance_scale=guidance_scale,83                    num_images_per_prompt=1,84                    negative_prompt=negative_prompt,85                    # height=height,86                    # width=width87        )88 89        # return first generated PIL image90        return out.images[0]91