Update app.py

app.py

@@ -52,10 +52,16 @@ def spherical_dist_loss(x, y):
     return (x - y).norm(dim=-1).div(2).arcsin().pow(2).mul(2)
     
 cc12m_model = hf_hub_download(repo_id="multimodalart/crowsonkb-v-diffusion-cc12m-1-cfg", filename="cc12m_1_cfg.pth")
+cc12m_small_model = hf_hub_download(repo_id="multimodalart/crowsonkb-v-diffusion-cc12m-1-cfg", filename="cc12m_1.pth")
 model = get_model('cc12m_1_cfg')()
 _, side_y, side_x = model.shape
 model.load_state_dict(torch.load(cc12m_model, map_location='cpu'))
 model = model.half().cuda().eval().requires_grad_(False)
+
+model_small = get_model('cc12m_1')()
+model_small.load_state_dict(torch.load(cc12m_model, map_location='cpu'))
+model_small = model.half().cuda().eval().requires_grad_(False)
+
 clip_model = clip.load(model.clip_model, jit=False, device='cuda')[0]
 clip_model.eval().requires_grad_(False)
 normalize = transforms.Normalize(mean=[0.48145466, 0.4578275, 0.40821073],
@@ -138,7 +144,7 @@ def run_all(prompt, steps, n_images, weight, clip_guided):
     else:
         extra_args = {'clip_embed': clip_embed}
         cond_fn_ = cond_fn
-        model_fn = make_cond_model_fn(model, cond_fn_)
+        model_fn = make_cond_model_fn(model_small, cond_fn_)
         outs = sampling.plms_sample(model_fn, x, step_list, extra_args)
     images_out = []
     for i, out in enumerate(outs):