Update app_2.py

app_2.py

@@ -104,41 +104,6 @@ download_models()
 
 
 
-@spaces.GPU()
-def infer(
-    prompt,
-    image,
-    do_rembg=False,
-    seed=42,
-    randomize_seed=False,
-    guidance_scale=3.0,
-    num_inference_steps=50,
-    reference_conditioning_scale=1.0,
-    negative_prompt="watermark, ugly, deformed, noisy, blurry, low contrast",
-    progress=gr.Progress(track_tqdm=True),
-):
-    # if do_rembg:
-    #     remove_bg_fn = lambda x: remove_bg(x, birefnet, transform_image, device)
-    # else:
-    #     remove_bg_fn = None
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-    images, preprocessed_image = run_pipeline(
-        pipe,
-        num_views=NUM_VIEWS,
-        text=prompt,
-        image=image,
-        height=HEIGHT,
-        width=WIDTH,
-        num_inference_steps=num_inference_steps,
-        guidance_scale=guidance_scale,
-        seed=seed,
-        remove_bg_fn=None,
-        reference_conditioning_scale=reference_conditioning_scale,
-        negative_prompt=negative_prompt,
-        device=device,
-    )
-    return images
 
 
 try:
@@ -381,6 +346,45 @@ def encode_prompt_pair(positive_prompt, negative_prompt):
 
     return c, uc
 
+@spaces.GPU(duration=60)
+@torch.inference_mode()
+def infer(
+    prompt,
+    image,
+    do_rembg=False,
+    seed=42,
+    randomize_seed=False,
+    guidance_scale=3.0,
+    num_inference_steps=50,
+    reference_conditioning_scale=1.0,
+    negative_prompt="watermark, ugly, deformed, noisy, blurry, low contrast",
+    progress=gr.Progress(track_tqdm=True),
+):
+    # if do_rembg:
+    #     remove_bg_fn = lambda x: remove_bg(x, birefnet, transform_image, device)
+    # else:
+    #     remove_bg_fn = None
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+        
+    images, preprocessed_image = run_pipeline(
+        pipe,
+        num_views=NUM_VIEWS,
+        text=prompt,
+        image=image,
+        height=HEIGHT,
+        width=WIDTH,
+        num_inference_steps=num_inference_steps,
+        guidance_scale=guidance_scale,
+        seed=seed,
+        remove_bg_fn=None,
+        reference_conditioning_scale=reference_conditioning_scale,
+        negative_prompt=negative_prompt,
+        device=device,
+    )
+    return images
+
+
 @spaces.GPU(duration=60)
 @torch.inference_mode()
 def pytorch2numpy(imgs, quant=True):
@@ -1228,7 +1232,8 @@ with block:
         example_quick_prompts.click(lambda x, y: ', '.join(y.split(', ')[:2] + [x[0]]), inputs=[example_quick_prompts, prompt], outputs=prompt, show_progress=False, queue=False)
         example_quick_subjects.click(lambda x: x[0], inputs=example_quick_subjects, outputs=prompt, show_progress=False, queue=False)
 
-        run_button.click(fn=infer,
+        run_button.click(
+            fn=infer,
             inputs=[
                 "high quality",
                 extracted_fg,