add inference

app.py

@@ -175,38 +175,40 @@ accelerator = accelerate.Accelerator()
 
 def generate_image_with_steps(prompt, negative_prompt, seed, width, height, guidance_scale, num_inference_steps):
     """Helper function to generate image with specific number of steps"""
-    prompts =  [prompt]
-    if num_inference_steps > 8:
-        sampler = UniPCSampler(pipe,model_closure=model_closure, steps=num_inference_steps, guidance_scale=guidance_scale,skip_type='time_uniform')
-    else:
-        sampler = UniPCSampler(pipe,model_closure=model_closure, steps=num_inference_steps, guidance_scale=guidance_scale)
-
-    c = prompts
-    uc = ['(worst quality:2), (low quality:2), (normal quality:2), bad anatomy, bad proportions, poorly drawn face, poorly drawn hands, missing fingers, extra limbs, blurry, pixelated, distorted, lowres, jpeg artifacts, watermark, signature, text, (deformed:1.5), (bad hands:1.3), overexposed, underexposed, censored, mutated, extra fingers, cloned face, bad eyes'] * len(c) if guidance_scale != 1.0 else None
-    shape = [4, width // 8, height // 8]
+    with torch.no_grad():
+        with precision_scope("cuda"):
+            prompts =  [prompt]
+            if num_inference_steps > 8:
+                sampler = UniPCSampler(pipe,model_closure=model_closure, steps=num_inference_steps, guidance_scale=guidance_scale,skip_type='time_uniform')
+            else:
+                sampler = UniPCSampler(pipe,model_closure=model_closure, steps=num_inference_steps, guidance_scale=guidance_scale)
+
+            c = prompts
+            uc = ['(worst quality:2), (low quality:2), (normal quality:2), bad anatomy, bad proportions, poorly drawn face, poorly drawn hands, missing fingers, extra limbs, blurry, pixelated, distorted, lowres, jpeg artifacts, watermark, signature, text, (deformed:1.5), (bad hands:1.3), overexposed, underexposed, censored, mutated, extra fingers, cloned face, bad eyes'] * len(c) if guidance_scale != 1.0 else None
+            shape = [4, width // 8, height // 8]
                     # if opt.method == "dpm_solver_v3":
                             # batch_size, shape, conditioning, x_T, unconditional_conditioning
-    samples, _ = sampler.sample(
-        conditioning=c,
-        batch_size=1,
-        shape=shape,
-        unconditional_conditioning=uc,
-        x_T=None,
-        start_free_u_step=6 if num_inference_steps == 8 else 4,
-        xl_preprocess_closure = prepare_sdxl_pipeline_step_parameter,
+            samples, _ = sampler.sample(
+                conditioning=c,
+                batch_size=1,
+                shape=shape,
+                unconditional_conditioning=uc,
+                x_T=None,
+                start_free_u_step=6 if num_inference_steps == 8 else 4 if num_inference_steps < 8 else None,
+                xl_preprocess_closure = prepare_sdxl_pipeline_step_parameter,
                         # npnet = npn_net,
-        use_corrector=True,
-    )
+                use_corrector=True,
+            )
 
-    x_samples = pipe.vae.decode(samples / pipe.vae.config.scaling_factor).sample
-    x_samples = torch.clamp((x_samples + 1.0) / 2.0, min=0.0, max=1.0)
-    x_samples = x_samples.cpu().permute(0, 2, 3, 1).numpy()
+            x_samples = pipe.vae.decode(samples / pipe.vae.config.scaling_factor).sample
+            x_samples = torch.clamp((x_samples + 1.0) / 2.0, min=0.0, max=1.0)
+            x_samples = x_samples.cpu().permute(0, 2, 3, 1).numpy()
 
-    x_image_torch = torch.from_numpy(x_samples).permute(0, 3, 1, 2) # need to pay attention
+            x_image_torch = torch.from_numpy(x_samples).permute(0, 3, 1, 2) # need to pay attention
 
-    for x_sample in x_image_torch:
-        x_sample = 255.0 * rearrange(x_sample.cpu().numpy(), "c h w -> h w c")
-        img = Image.fromarray(x_sample.astype(np.uint8))
+            for x_sample in x_image_torch:
+                x_sample = 255.0 * rearrange(x_sample.cpu().numpy(), "c h w -> h w c")
+                img = Image.fromarray(x_sample.astype(np.uint8))
     return img
 
 @spaces.GPU #[uncomment to use ZeroGPU]