Fix ui and add apex

app.py

@@ -13,7 +13,6 @@ model_repo_id = "michellemoorre/var-test"
 pipe = TVARPipeline.from_pretrained(model_repo_id, device=device)
 
 MAX_SEED = np.iinfo(np.int32).max
-MAX_IMAGE_SIZE = 1024
 
 
 @spaces.GPU(duration=65)
@@ -25,8 +24,9 @@ def infer(
     guidance_scale=4.0,
     top_k=450,
     top_p=0.95,
-    re=False,
+    re=True,
     re_max_depth=10,
+    re_start_iter=2,
     progress=gr.Progress(track_tqdm=True),
 ):
     if randomize_seed:
@@ -39,6 +39,8 @@ def infer(
         top_p=top_p,
         top_k=top_k,
         re=re,
+        re_max_depth=re_max_depth,
+        re_start_iter=re_start_iter,
         g_seed=seed,
     )[0]
 
@@ -73,6 +75,23 @@ with gr.Blocks(css=css) as demo:
 
         result = gr.Image(label="Result", show_label=False)
 
+        seed = gr.Number(
+            label="Seed",
+            minimum=0,
+            maximum=MAX_SEED,
+            value=0,
+        )
+
+        randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+
+        guidance_scale = gr.Slider(
+            label="Guidance scale",
+            minimum=0.0,
+            maximum=7.5,
+            step=0.5,
+            value=4.,
+        )
+
         with gr.Accordion("Advanced Settings", open=False):
             negative_prompt = gr.Text(
                 label="Negative prompt",
@@ -81,30 +100,12 @@ with gr.Blocks(css=css) as demo:
                 visible=True,
             )
 
-            seed = gr.Slider(
-                label="Seed",
-                minimum=0,
-                maximum=MAX_SEED,
-                step=1,
-                value=0,
-            )
-
-            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-
-            with gr.Row():
-                guidance_scale = gr.Slider(
-                    label="Guidance scale",
-                    minimum=0.0,
-                    maximum=7.5,
-                    step=0.1,
-                    value=4.5,
-                )
             with gr.Row():
                 top_k = gr.Slider(
                     label="Sampling top k",
-                    minimum=1,
+                    minimum=10,
                     maximum=1000,
-                    step=10,
+                    step=20,
                     value=450,
                 )
                 top_p = gr.Slider(
@@ -114,15 +115,23 @@ with gr.Blocks(css=css) as demo:
                     step=0.05,
                     value=0.95,
                 )
+
+            re = gr.Checkbox(label="Rejection Sampling (RE)", value=True)
             with gr.Row():
-                re = gr.Checkbox(label="Rejection Sampling", value=False)
                 re_max_depth = gr.Slider(
-                    label="Rejection Sampling Depth",
+                    label="RE Depth",
                     minimum=0,
                     maximum=20,
-                    step=1,
+                    step=4,
                     value=10,
                 )
+                re_start_iter = gr.Slider(
+                    label="RE Start Scale",
+                    minimum=0,
+                    maximum=9,
+                    step=1,
+                    value=2,
+                )
 
         gr.Examples(examples=examples, inputs=[prompt], outputs=[result, seed], fn=infer, cache_examples=True)# cache_mode="lazy")
     gr.on(
@@ -138,6 +147,7 @@ with gr.Blocks(css=css) as demo:
             top_p,
             re,
             re_max_depth,
+            re_start_iter,
         ],
         outputs=[result, seed],
     )

models/pipeline.py

@@ -96,6 +96,7 @@ class TVARPipeline:
         more_smooth=False,
         re=False,
         re_max_depth=10,
+        re_start_iter=2,
         return_pil=True,
         encoded_prompt = None,
         encoded_null_prompt = None,
@@ -180,7 +181,7 @@ class TVARPipeline:
             idx_Bl = sample_with_top_k_top_p_(
                 logits_BlV, rng=rng, top_k=top_k, top_p=top_p, num_samples=1
             )[:, :, 0]
-            if re:
+            if re and si >= re_start_iter:
                 selected_logits = torch.gather(logits_BlV, -1, idx_Bl.unsqueeze(-1))[:, :, 0]
                 mx = selected_logits.sum(dim=-1)[:, None]
                 for _ in range(re_max_depth):

requirements.txt

@@ -1,4 +1,5 @@
 torch==2.5.0
+apex==0.9.10dev
 decord==0.6.0
 einops==0.8.0
 huggingface_hub==0.26.1