diff --git a/entry.py b/entry.py
index d7f773d1..9cdd9245 100644
--- a/entry.py
+++ b/entry.py
@@ -98,60 +98,59 @@ model.load_state_dict(safetensors.torch.load_file('./sd_xl_base_1.0.safetensors'
 
 # model.conditioner.cuda()
 
-model.conditioner.embedders[0].device = 'cpu'
-model.conditioner.embedders[1].device = 'cpu'
-
-value_dict = {
-    "prompt": "a handsome man in forest", "negative_prompt": "ugly, bad", "orig_height": 1024, "orig_width": 1024,
-    "crop_coords_top": 0, "crop_coords_left": 0, "target_height": 1024, "target_width": 1024, "aesthetic_score": 7.5,
-    "negative_aesthetic_score": 2.0,
-}
-
-batch, batch_uc = get_batch(
-    get_unique_embedder_keys_from_conditioner(model.conditioner),
-    value_dict,
-    [1],
-)
-
-c, uc = model.conditioner.get_unconditional_conditioning(
-    batch,
-    batch_uc=batch_uc)
-# model.conditioner.cpu()
-
-c = {a: b.to(torch.float16) for a, b in c.items()}
-uc = {a: b.to(torch.float16) for a, b in uc.items()}
-
-torch.cuda.empty_cache()
-torch.cuda.ipc_collect()
-
-shape = (1, 4, 128, 128)
-randn = torch.randn(shape).to(torch.float16).cuda()
-
-
-def denoiser(input, sigma, c):
-    return model.denoiser(model.model, input, sigma, c)
-
-
 with torch.no_grad():
-    model.model.to(torch.float16).cuda()
-    model.denoiser.to(torch.float16).cuda()
-    samples_z = sampler(denoiser, randn, cond=c, uc=uc)
-    model.model.cpu()
-    model.denoiser.cpu()
 
-torch.cuda.empty_cache()
-torch.cuda.ipc_collect()
+    model.conditioner.embedders[0].device = 'cpu'
+    model.conditioner.embedders[1].device = 'cpu'
 
-a = 0
+    value_dict = {
+        "prompt": "a handsome man in forest", "negative_prompt": "ugly, bad", "orig_height": 1024, "orig_width": 1024,
+        "crop_coords_top": 0, "crop_coords_left": 0, "target_height": 1024, "target_width": 1024, "aesthetic_score": 7.5,
+        "negative_aesthetic_score": 2.0,
+    }
+
+    batch, batch_uc = get_batch(
+        get_unique_embedder_keys_from_conditioner(model.conditioner),
+        value_dict,
+        [1],
+    )
+
+    c, uc = model.conditioner.get_unconditional_conditioning(
+        batch,
+        batch_uc=batch_uc)
+    # model.conditioner.cpu()
+
+    c = {a: b.to(torch.float16) for a, b in c.items()}
+    uc = {a: b.to(torch.float16) for a, b in uc.items()}
+
+    torch.cuda.empty_cache()
+    torch.cuda.ipc_collect()
+
+    shape = (1, 4, 128, 128)
+    randn = torch.randn(shape).to(torch.float16).cuda()
 
 
-with torch.no_grad():
-    model.first_stage_model.cuda()
-    samples_x = model.decode_first_stage(samples_z)
-    samples = torch.clamp((samples_x + 1.0) / 2.0, min=0.0, max=1.0)
-    model.first_stage_model.cpu()
+    def denoiser(input, sigma, c):
+        return model.denoiser(model.model, input, sigma, c)
 
-import cv2
-samples = einops.rearrange(samples, 'b c h w -> b h w c')[0] * 127.5 + 127.5
-samples = samples.cpu().numpy().clip(0, 255).astype(np.uint8)
-cv2.imwrite('img.png', samples)
+
+    with torch.no_grad():
+        model.model.to(torch.float16).cuda()
+        model.denoiser.to(torch.float16).cuda()
+        samples_z = sampler(denoiser, randn, cond=c, uc=uc)
+        model.model.cpu()
+        model.denoiser.cpu()
+
+    torch.cuda.empty_cache()
+    torch.cuda.ipc_collect()
+
+    with torch.no_grad():
+        model.first_stage_model.to(torch.float16).cuda()
+        samples_x = model.decode_first_stage(samples_z)
+        samples = torch.clamp((samples_x + 1.0) / 2.0, min=0.0, max=1.0)
+        model.first_stage_model.cpu()
+
+    import cv2
+    samples = einops.rearrange(samples, 'b c h w -> b h w c')[0] * 127.5 + 127.5
+    samples = samples.cpu().numpy().clip(0, 255).astype(np.uint8)
+    cv2.imwrite('img.png', samples)