Update app.py

app.py

@@ -10,7 +10,7 @@ import numpy as np
 from PIL import Image, ImageDraw, ImageFont
 from functools import partial
 import math
-from utils import compute_ca_loss
+from utils import compute_loco_v2
 from gradio import processing_utils
 from typing import Optional
 
@@ -95,10 +95,11 @@ def inference(device, unet, vae, tokenizer, text_encoder, prompt, bboxes, object
         generator=generator,
     ).to(device)
 
-    noise_scheduler = LMSDiscreteScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000)
+    # noise_scheduler = LMSDiscreteScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000)
+    noise_scheduler = DDIMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", clip_sample=False, set_alpha_to_one=False)
 
     # generator = torch.Generator("cuda").manual_seed(1024)
-    noise_scheduler.set_timesteps(51)
+    noise_scheduler.set_timesteps(50)
 
     latents = latents * noise_scheduler.init_noise_sigma
 
@@ -119,14 +120,14 @@ def inference(device, unet, vae, tokenizer, text_encoder, prompt, bboxes, object
 
             # update latents with guidence from gaussian blob
 
-            loss = compute_ca_loss(attn_map_integrated_mid, attn_map_integrated_up, bboxes=bboxes,
+            loss = compute_loco_v2(attn_map_integrated_mid, attn_map_integrated_up, bboxes=bboxes,
                                    object_positions=object_positions) * loss_scale
 
-            print(loss.item() / loss_scale)
+            # print(loss.item() / loss_scale)
 
             grad_cond = torch.autograd.grad(loss.requires_grad_(True), [latents])[0]
 
-            latents = latents - grad_cond * noise_scheduler.sigmas[index] ** 2
+            latents = latents - grad_cond 
             iteration += 1
             torch.cuda.empty_cache()
         torch.cuda.empty_cache()