ciCic/Image-Encoder

1import os2import torch3import gradio as gr4 5from diffusers import AutoencoderTiny6from torchvision.transforms.functional import to_pil_image, center_crop, resize, to_tensor7 8device = 'cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu'9vae = None10 11 12def get_model():13    global vae14 15    if vae is None:16        model_id = "madebyollin/taesd"17        vae = AutoencoderTiny.from_pretrained(model_id, safetensors=True).to(device)18    return vae19 20 21@torch.no_grad()22def encode(image):23    vae = get_model()24 25    DIM = [512]26    processed = center_crop(resize(image, DIM), DIM)27    tensor = to_tensor(processed).unsqueeze(0).to(device)28    latents = vae.encoder(tensor)29    scaled = vae.scale_latents(latents).mul_(255).round_().byte()30    return to_pil_image(scaled[0])31 32 33path = 'https://huggingface.co/buckets/ciCic/demo-purposes/resolve/images'34astronaut = f"{path}/6.png"35 36 37def app():38    return gr.Interface(encode,39                        gr.Image(type="pil",40                                 label='512x512',41                                 value=astronaut),42                        gr.Image(type="pil",43                                 image_mode="RGBA",44                                 label='64x64',45                                 height=256,46                                 width=25647                                 ),48                        examples=[49                            astronaut,50                            f"{path}/7.png",51                            f"{path}/34.png"52                        ], flagging_mode='never', title='Image Encoder')53 54 55if __name__ == "__main__":56    print("LAUNCHING")57    app().launch(server_name="0.0.0.0", server_port=7860, share=True)58