import torch
from chat_anything.face_generator.pipelines.lpw_stable_diffusion import StableDiffusionLongPromptWeightingPipeline 

@torch.no_grad()
def generate(pipe, prompt, negative_prompt, **generating_conf):
    pipe_longprompt = StableDiffusionLongPromptWeightingPipeline(
        unet=pipe.unet,
        text_encoder=pipe.text_encoder,
        vae=pipe.vae,
        tokenizer=pipe.tokenizer,
        scheduler=pipe.scheduler, 
        safety_checker=None,
        feature_extractor=None,
    )
    print('generating: ', prompt)
    print('using negative prompt: ', negative_prompt)
    embeds = pipe_longprompt._encode_prompt(prompt=prompt, negative_prompt=negative_prompt, device=pipe.device, num_images_per_prompt=1, do_classifier_free_guidance=generating_conf['guidance_scale']>1,)
    negative_prompt_embeds, prompt_embeds = embeds.split(embeds.shape[0]//2)
    pipe_out = pipe(
        prompt_embeds=prompt_embeds,
        negative_prompt_embeds=negative_prompt_embeds,
        **generating_conf,
    )
    return pipe_out

if __name__ == '__main__':
    from diffusers.pipelines import StableDiffusionPipeline
    import argparse
    def main():
        parser = argparse.ArgumentParser()
        parser.add_argument(
            '--prompts',type=str,default=['starry night','Impression Sunrise, drawn by Claude Monet'], nargs='*'
        )
        
        args = parser.parse_args()
        prompts = args.prompts
        print(f'generating {prompts}')
        model_id = 'pretrained_model/sd-v1-4'
        pipe = StableDiffusionPipeline.from_pretrained(model_id,).to('cuda')
        images = pipe(prompts).images
        for i, image in enumerate(images):
            image.save(f'{prompts[i]}_{i}.png')

    main()