Update app.py

app.py

@@ -8,10 +8,15 @@ from sklearn.svm import LinearSVC
 from sklearn import preprocessing
 import pandas as pd
 
+from diffusers import LCMScheduler
+from diffusers.models import ImageProjection
+from patch_sdxl import SDEmb
+import torch
+import spaces
+
 import random
 import time
 
-import replicate
 import torch
 from urllib.request import urlopen
 
@@ -24,11 +29,49 @@ prompt_list = [p for p in list(set(
 
 start_time = time.time()
 
-# TODO add to state instead of shared across all
-glob_idx = 0
+####################### Setup Model
+model_id = "stabilityai/stable-diffusion-xl-base-1.0"
+lcm_lora_id = "latent-consistency/lcm-lora-sdxl"
+pipe = SDEmb.from_pretrained(model_id, variant="fp16")
+pipe.load_lora_weights(lcm_lora_id)
+pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
+pipe.to(device='cuda', dtype=torch.float16)
+pipe.load_ip_adapter("h94/IP-Adapter", subfolder="sdxl_models", weight_name="ip-adapter_sdxl.bin")
+output_hidden_state = False
+#######################
+
+@spaces.GPU
+def predict(
+        prompt,
+        im_emb=None,
+    ):
+    """Run a single prediction on the model"""
+    with torch.no_grad():
+        if im_emb == None:
+            im_emb = torch.zeros(1, 1280, dtype=torch.float16, device='cuda')
+        else:
+            im_emb = torch.tensor([float(i) for i in im_emb.split(', ')]).unsqueeze(0).to(dtype=torch.float16).to('cuda')
+        image = pipe(
+            prompt=prompt,
+            ip_adapter_emb=im_emb,
+            height=1024,
+            width=1024,                                                                                                             
+            num_inference_steps=8,
+            guidance_scale=0,
+            ).images[0]
+        im_emb, _ = pipe.encode_image(
+                image, 'cuda', 1, output_hidden_state
+            )
+        return image, im_emb.to(DEVICE)
+
+
+
 
 
-deployment = replicate.deployments.get("rynmurdock/zahir-deployment")
+
+
+# TODO add to state instead of shared across all
+glob_idx = 0
 
 def next_image(embs, ys, calibrate_prompts):
     global glob_idx
@@ -46,26 +89,11 @@ def next_image(embs, ys, calibrate_prompts):
             print('######### Calibrating with sample prompts #########')
             prompt = calibrate_prompts.pop(0)
             print(prompt)
-
-            prediction = deployment.predictions.create(
-              input={"prompt": prompt,}
-            )
-            prediction.wait()
-            output = prediction.output
-
-            # output = replicate.run(
-            # "rynmurdock/zahir:42c58addd49ab57f1e309f0b9a0f271f483bbef0470758757c623648fe989e42",
-            # input={"prompt": prompt,}
-            # )
-            
-            response = requests.get(output['file1'])
-            image = Image.open(BytesIO(response.content))
-
-            embs.append(torch.tensor([float(i) for i in urlopen(output['file2']).read().decode('utf-8').split(', ')]).unsqueeze(0))
+            image, img_emb = predict(prompt)
+            embs.append(img_emb)
             return image, embs, ys, calibrate_prompts
         else:
             print('######### Roaming #########')
-
             # sample only as many negatives as there are positives
             indices = range(len(ys))
             pos_indices = [i for i in indices if ys[i] == 1]
@@ -93,28 +121,8 @@ def next_image(embs, ys, calibrate_prompts):
             im_emb = w * lin_class.coef_.to(device=DEVICE, dtype=torch.float16)
             prompt= 'an image' if glob_idx % 2 == 0 else rng_prompt
             print(prompt)
-
-            im_emb_st = str(im_emb[0].cpu().detach().tolist())[1:-1]
-            
-            prediction = deployment.predictions.create(
-              input={"prompt": prompt, 'im_emb': im_emb_st}
-            )
-            prediction.wait()
-            output = prediction.output
-            
-            # output = replicate.run(
-            # "rynmurdock/zahir:42c58addd49ab57f1e309f0b9a0f271f483bbef0470758757c623648fe989e42",
-            # input={"prompt": prompt, 'im_emb': im_emb_st}
-            # )
-            
-            response = requests.get(output['file1'])
-            image = Image.open(BytesIO(response.content))
-
-
-            im_emb = torch.tensor([float(i) for i in urlopen(output['file2']).read().decode('utf-8').split(', ')]).unsqueeze(0)
+            image, im_emb = predict(prompt, img_emb)
             embs.append(im_emb)
-
-            torch.save(lin_class.coef_, f'./{start_time}.pt')
             return image, embs, ys, calibrate_prompts
 
 
@@ -195,6 +203,6 @@ with gr.Blocks(css=css) as demo:
                  [b4, embs, ys, calibrate_prompts],
                  [b1, b2, b3, b4, img, embs, ys, calibrate_prompts])
     with gr.Row():
-        html = gr.HTML('''<div style='text-align:center; font-size:32'>You will callibrate for several prompts and then roam.</ div>''')
+        html = gr.HTML('''<div style='text-align:center; font-size:32'>You will calibrate for several prompts and then roam.</ div>''')
 
 demo.launch()  # Share your demo with just 1 extra parameter 🚀