Spaces:
Sleeping
Sleeping
File size: 1,738 Bytes
761d68b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 |
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 |