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