| 668 | |
| 669 | |
| 670 | class DepthText2Image: |
| 671 | def __init__(self, device): |
| 672 | print(f"Initializing DepthText2Image to {device}") |
| 673 | self.torch_dtype = torch.float16 if 'cuda' in device else torch.float32 |
| 674 | self.controlnet = ControlNetModel.from_pretrained( |
| 675 | "fusing/stable-diffusion-v1-5-controlnet-depth", torch_dtype=self.torch_dtype) |
| 676 | self.pipe = StableDiffusionControlNetPipeline.from_pretrained( |
| 677 | "runwayml/stable-diffusion-v1-5", controlnet=self.controlnet, safety_checker=StableDiffusionSafetyChecker.from_pretrained('CompVis/stable-diffusion-safety-checker'), |
| 678 | torch_dtype=self.torch_dtype) |
| 679 | self.pipe.scheduler = UniPCMultistepScheduler.from_config(self.pipe.scheduler.config) |
| 680 | self.pipe.to(device) |
| 681 | self.seed = -1 |
| 682 | self.a_prompt = 'best quality, extremely detailed' |
| 683 | self.n_prompt = 'longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit,' \ |
| 684 | ' fewer digits, cropped, worst quality, low quality' |
| 685 | |
| 686 | @prompts(name="Generate Image Condition On Depth", |
| 687 | description="useful when you want to generate a new real image from both the user description and depth image. " |
| 688 | "like: generate a real image of a object or something from this depth image, " |
| 689 | "or generate a new real image of a object or something from the depth map. " |
| 690 | "The input to this tool should be a comma separated string of two, " |
| 691 | "representing the image_path and the user description") |
| 692 | def inference(self, inputs): |
| 693 | image_path, instruct_text = inputs.split(",")[0], ','.join(inputs.split(',')[1:]) |
| 694 | image = Image.open(image_path) |
| 695 | self.seed = random.randint(0, 65535) |
| 696 | seed_everything(self.seed) |
| 697 | prompt = f'{instruct_text}, {self.a_prompt}' |
| 698 | image = self.pipe(prompt, image, num_inference_steps=20, eta=0.0, negative_prompt=self.n_prompt, |
| 699 | guidance_scale=9.0).images[0] |
| 700 | updated_image_path = get_new_image_name(image_path, func_name="depth2image") |
| 701 | image.save(updated_image_path) |
| 702 | print(f"\nProcessed DepthText2Image, Input Depth: {image_path}, Input Text: {instruct_text}, " |
| 703 | f"Output Image: {updated_image_path}") |
| 704 | return updated_image_path |
| 705 | |
| 706 | |
| 707 | class Image2Normal: |
nothing calls this directly
no outgoing calls
no test coverage detected