Spaces:
Sleeping
Sleeping
| from diffusers import StableDiffusionPipeline,StableDiffusionControlNetPipeline, ControlNetModel, UniPCMultistepScheduler | |
| from diffusers.utils import load_image | |
| import torch | |
| def generate_image(model_name,input_text): | |
| pipe = StableDiffusionPipeline.from_pretrained(model_name, torch_dtype=torch.float16) | |
| # pipe = pipe.to("cuda") | |
| prompt = input_text | |
| image = pipe(prompt).images[0] | |
| image.save("testo.png") | |
| return image | |
| def generate_controlnet_image(model_name,algorithm,input_image,input_text): | |
| mask_image = generate_mask(input_image,algorithm) | |
| base_model_path = model_name | |
| controlnet_path = "lllyasviel/control_v11p_sd15_inpaint" | |
| controlnet = ControlNetModel.from_pretrained(controlnet_path, torch_dtype=torch.float16) | |
| pipe = StableDiffusionControlNetPipeline.from_pretrained( | |
| base_model_path, controlnet=controlnet, torch_dtype=torch.float16 | |
| ) | |
| # speed up diffusion process with faster scheduler and memory optimization | |
| pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config) | |
| # remove following line if xformers is not installed or when using Torch 2.0. | |
| pipe.enable_xformers_memory_efficient_attention() | |
| # memory optimization. | |
| pipe.enable_model_cpu_offload() | |
| control_image = load_image(mask_image) | |
| prompt = "pale golden rod circle with old lace background" | |
| # generate image | |
| generator = torch.manual_seed(0) | |
| image = pipe( | |
| prompt, num_inference_steps=20, generator=generator, image=control_image | |
| ).images[0] | |
| image.save("./output.png") | |
| return mask_image | |
| def generate_video(model_name,input_image,input_text): | |
| return input_image | |
| def generate_mask(image,algorithm): | |
| pass |