diffusers/tools

1#!/usr/bin/env python32from huggingface_hub import HfApi3import torch4 5import requests6from PIL import Image7 8from diffusers import DDIMScheduler, StableDiffusionPix2PixZeroPipeline9from diffusers.schedulers.scheduling_ddim_inverse import DDIMInverseScheduler10from transformers import BlipForConditionalGeneration, BlipProcessor11 12api = HfApi()13img_url = "https://github.com/pix2pixzero/pix2pix-zero/raw/main/assets/test_images/cats/cat_6.png"14raw_image = Image.open(requests.get(img_url, stream=True).raw).convert('RGB').resize((512, 512))15 16processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")17model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base", torch_dtype=torch.float16, low_cpu_mem_usage=True)18 19model_ckpt = "CompVis/stable-diffusion-v1-4"20pipeline = StableDiffusionPix2PixZeroPipeline.from_pretrained(21    model_ckpt, caption_generator=model, caption_processor=processor, torch_dtype=torch.float16, safety_checker=None,22)23pipeline.enable_model_cpu_offload()24 25caption = pipeline.generate_caption(raw_image)26 27pipeline.scheduler = DDIMScheduler.from_config(pipeline.scheduler.config)28pipeline.inverse_scheduler = DDIMInverseScheduler.from_config(pipeline.scheduler.config)29 30print(caption)31 32generator = torch.manual_seed(0)33inv_latents = pipeline.invert(caption, image=raw_image, generator=generator).latents34 35source_prompts = 4 * ["a cat sitting on the street", "a cat playing in the field", "a face of a cat"]36target_prompts = 4 * ["a dog sitting on the street", "a dog playing in the field", "a face of a dog"]37 38source_embeds = pipeline.get_embeds(source_prompts, batch_size=2)39target_embeds = pipeline.get_embeds(target_prompts, batch_size=2)40 41 42image = pipeline(43    caption,44    source_embeds=source_embeds,45    target_embeds=target_embeds,46    num_inference_steps=50,47    cross_attention_guidance_amount=0.15,48    generator=generator,49    latents=inv_latents,50    negative_prompt=caption,51).images[0]52 53path = "/home/patrick_huggingface_co/images/aa.png"54image.save(path)55 56api.upload_file(57    path_or_fileobj=path,58    path_in_repo=path.split("/")[-1],59    repo_id="patrickvonplaten/images",60    repo_type="dataset",61)62print("https://huggingface.co/datasets/patrickvonplaten/images/blob/main/aa.png")63