using torch.compile over sfast for compatibility; other smol changes

app.py

@@ -1,23 +1,22 @@
 DEVICE = 'cuda'
 
-from sfast.compilers.diffusion_pipeline_compiler import (compile,
-                                                         CompilationConfig)
-config = CompilationConfig.Default()
-
 import gradio as gr
 import numpy as np
 from sklearn.svm import LinearSVC
 from sklearn import preprocessing
 import pandas as pd
 
-from diffusers import LCMScheduler, AutoencoderTiny, EulerDiscreteScheduler, UNet2DConditionModel, AutoPipelineForText2Image
+from diffusers import LCMScheduler, AutoencoderTiny, EulerDiscreteScheduler, UNet2DConditionModel, AutoPipelineForText2Image, DiffusionPipeline
 from diffusers.models import ImageProjection
 import torch
 
+torch.set_float32_matmul_precision('high')
+
 import random
 import time
 
-import torch
+# TODO put back
+import spaces
 from urllib.request import urlopen
 
 from PIL import Image
@@ -27,7 +26,6 @@ from io import BytesIO, StringIO
 from transformers import CLIPVisionModelWithProjection
 from huggingface_hub import hf_hub_download
 from safetensors.torch import load_file
-#import spaces
 
 prompt_list = [p for p in list(set(
                 pd.read_csv('./twitter_prompts.csv').iloc[:, 1].tolist())) if type(p) == str]
@@ -35,6 +33,7 @@ prompt_list = [p for p in list(set(
 start_time = time.time()
 
 ####################### Setup Model
+
 model_id = "stabilityai/stable-diffusion-xl-base-1.0"
 sdxl_lightening = "ByteDance/SDXL-Lightning"
 ckpt = "sdxl_lightning_2step_unet.safetensors"
@@ -46,20 +45,20 @@ pipe = AutoPipelineForText2Image.from_pretrained(model_id, unet=unet, torch_dtyp
 pipe.unet._load_ip_adapter_weights(torch.load(hf_hub_download('h94/IP-Adapter', 'sdxl_models/ip-adapter_sdxl_vit-h.bin')))
 pipe.load_ip_adapter("h94/IP-Adapter", subfolder="sdxl_models", weight_name="ip-adapter_sdxl_vit-h.bin")
 pipe.register_modules(image_encoder = image_encoder)
+pipe.set_ip_adapter_scale(0.8)
 
 pipe.vae = AutoencoderTiny.from_pretrained("madebyollin/taesdxl", torch_dtype=torch.float16)
 pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing")
+
 pipe.to(device=DEVICE)
-pipe = compile(pipe, config=config)
 
-image = pipe(prompt_embeds=torch.zeros(1, 1, 2048, dtype=torch.float16, device=DEVICE),
-                pooled_prompt_embeds=torch.zeros(1, 1280, dtype=torch.float16, device=DEVICE),
-                ip_adapter_image_embeds=[torch.zeros(1, 1, 1024, dtype=torch.float16, device=DEVICE)],
-                height=1024,
-                width=1024,                                                                                                             
-                num_inference_steps=2,
-                guidance_scale=0,
-            ).images[0]
+
+# TODO put back
+@spaces.GPU
+def compile_em():
+    pipe.unet = torch.compile(pipe.unet)
+    pipe.vae = torch.compile(pipe.vae)
+    autoencoder.model.forward = torch.compile(autoencoder.model.forward, backend='inductor', dynamic=True)
 
 
 output_hidden_state = False
@@ -76,7 +75,6 @@ class BottleneckT5Autoencoder:
         self.tokenizer = AutoTokenizer.from_pretrained(model_path, model_max_length=512, torch_dtype=torch.bfloat16)
         self.model = AutoModelForCausalLM.from_pretrained(model_path, trust_remote_code=True).to(self.device)
         self.model.eval()
-        # self.model = torch.compile(self.model)
 
 
     def embed(self, text: str) -> torch.FloatTensor:
@@ -88,7 +86,7 @@ class BottleneckT5Autoencoder:
             encode_only=True,
         )
 
-    def generate_from_latent(self, latent: torch.FloatTensor, max_length=512, temperature=1., top_p=.8, length_penalty=10, min_new_tokens=30) -> str:
+    def generate_from_latent(self, latent: torch.FloatTensor, max_length=512, temperature=1., top_p=.8, min_new_tokens=30) -> str:
         dummy_text = '.'
         dummy = self.embed(dummy_text)
         perturb_vector = latent - dummy
@@ -101,7 +99,6 @@ class BottleneckT5Autoencoder:
             temperature=temperature,
             top_p=top_p,
             num_return_sequences=1,
-            length_penalty=length_penalty,
             min_new_tokens=min_new_tokens,
             # num_beams=8,
         )
@@ -109,8 +106,11 @@ class BottleneckT5Autoencoder:
 
 autoencoder = BottleneckT5Autoencoder(model_path='thesephist/contra-bottleneck-t5-xl-wikipedia')
 
+compile_em()
 #######################
 
+# TODO put back
+@spaces.GPU
 def generate(prompt, in_embs=None,):
   if prompt != '':
     print(prompt)
@@ -119,12 +119,12 @@ def generate(prompt, in_embs=None,):
   else:
     print('From embeds.')
   in_embs = in_embs / in_embs.abs().max() * .15
-  text = autoencoder.generate_from_latent(in_embs.to('cuda').to(dtype=torch.bfloat16), temperature=.3, top_p=.99, min_new_tokens=5)
+  text = autoencoder.generate_from_latent(in_embs.to('cuda').to(dtype=torch.bfloat16), temperature=.8, top_p=.94, min_new_tokens=5)
   return text, in_embs.to('cpu')
 
 
-
-#@spaces.GPU
+# TODO put back
+@spaces.GPU
 def predict(
         prompt,
         im_emb=None,
@@ -145,6 +145,7 @@ def predict(
                 width=1024,                                                                                                             
                 num_inference_steps=2,
                 guidance_scale=0,
+#                timesteps=[800],
             ).images[0]
         else:
             image = pipe(
@@ -154,6 +155,7 @@ def predict(
                 width=1024,                                                                                                             
                 num_inference_steps=2,
                 guidance_scale=0,
+#                timesteps=[800],
             ).images[0]
         im_emb, _ = pipe.encode_image(
                 image, DEVICE, 1, output_hidden_state
@@ -232,7 +234,7 @@ def next_image(embs, img_embs, ys, calibrate_prompts):
             rng_prompt = random.choice(prompt_list)
             w = 1.4# if len(embs) % 2 == 0 else 0
             
-            prompt= '' if glob_idx % 2 == 0 else rng_prompt
+            prompt= '' if glob_idx % 3 == 0 else rng_prompt
             prompt, _ = generate(prompt, in_embs=im_s)
             print(prompt)
             im_emb = autoencoder.embed(prompt)

requirements.txt

@@ -8,4 +8,3 @@ diffusers
 accelerate
 transformers
 peft
-https://github.com/chengzeyi/stable-fast/releases/download/v1.0.4/stable_fast-1.0.4+torch220cu121-cp310-cp310-manylinux2014_x86_64.whl