update app

app.py

@@ -21,39 +21,13 @@ from typing import List
 # Adjust the import path if your file is located elsewhere
 from src.smc.inference import infer_pretrained, infer_smc_grad, PretrainedInferenceConfig, SMCGradInferenceConfig
 
-# Global constants (adjust if needed)
-MAX_SEED = 2 ** 32 - 1
-MAX_IMAGE_SIZE = 1024
-DEVICE = "cpu"
-
-
-# Sensible defaults (change to match your model constraints)
-DEFAULTS = {
-    "resolution": 512,
-    "pretrained_steps": 20,
-    "pretrained_CFG": 7.5,
-    "pretrained_num_batches": 1,
-
-    "smc_steps": 20,
-    "smc_CFG": 7.5,
-    "smc_num_batches": 1,
-    "smc_num_particles": 4,
-    "smc_ess_threshold": 0.5,
-    "smc_partial_resampling": True,
-    "smc_resample_frequency": 5,
-    "smc_kl_weight": 0.1,
-    "smc_lambda_tempering": False,
-    "smc_lambda_one_at": 0.5,
-    "smc_phi": 1,
-    "smc_tau": 0.1,
-}
+DEVICE = "cuda"
 
 examples = [
     "A dreamy Monet-style landscape with soft brush strokes",
     "Vibrant city street at dawn in impressionist style",
 ]
 
-
 def _format_inference_output(out) -> str:
     """Return a short summary string for the UI"""
     if out is None:
@@ -71,34 +45,28 @@ def run_inference_all(
 
     # Pretrained method controls
     pretrained_negative_prompt,
-    pretrained_resolution,
     pretrained_CFG,
     pretrained_steps,
-    pretrained_num_batches,
-    pretrained_device,
 
     # SMC-grad method controls
-    smc_negative_prompt,
-    smc_resolution,
-    smc_CFG,
-    smc_steps,
-    smc_num_batches,
-    smc_num_particles,
-    smc_ess_threshold,
-    smc_partial_resampling,
-    smc_resample_frequency,
-    smc_kl_weight,
-    smc_lambda_tempering,
-    smc_lambda_one_at,
-    smc_use_continuous_formulation,
-    smc_phi,
-    smc_tau,
-    smc_proposal_type,
+    smc_grad_negative_prompt,
+    smc_grad_CFG,
+    smc_grad_steps,
+    smc_grad_num_particles,
+    smc_grad_ess_threshold,
+    smc_grad_partial_resampling,
+    smc_grad_resample_frequency,
+    smc_grad_kl_weight,
+    smc_grad_lambda_tempering,
+    smc_grad_lambda_one_at,
+    smc_grad_use_continuous_formulation,
+    smc_grad_phi,
+    smc_grad_tau,
 ):
     """Wrapper that runs both inference methods and returns UI-friendly outputs.
 
     Returns:
-      pretrained_images, pretrained_info, smc_images, smc_info
+      pretrained_images, pretrained_info, smc_grad_images, smc_grad_info
     """
     # --- Pretrained ---
     pretrained_output = None
@@ -107,12 +75,10 @@ def run_inference_all(
         pretrained_cfg = PretrainedInferenceConfig(
             prompt=prompt,
             negative_prompt=pretrained_negative_prompt or "",
-            resolution=int(pretrained_resolution),
             CFG=float(pretrained_CFG),
             steps=int(pretrained_steps),
-            num_batches=int(pretrained_num_batches),
         )
-        pretrained_output = infer_pretrained(pretrained_cfg, device=pretrained_device)
+        pretrained_output = infer_pretrained(pretrained_cfg, device=DEVICE)
         pretrained_images = pretrained_output.images
     except Exception as e:
         pretrained_images = []
@@ -121,44 +87,41 @@ def run_inference_all(
         pretrained_images = [pretrained_error]
 
     # --- SMC-grad ---
-    smc_output = None
-    smc_images = []
+    smc_grad_output = None
+    smc_grad_images = []
     try:
-        smc_cfg = SMCGradInferenceConfig(
+        smc_grad_cfg = SMCGradInferenceConfig(
             prompt=prompt,
-            negative_prompt=smc_negative_prompt or "",
-            ess_threshold=float(smc_ess_threshold),
-            partial_resampling=bool(smc_partial_resampling),
-            resample_frequency=int(smc_resample_frequency),
-            resolution=int(smc_resolution),
-            CFG=float(smc_CFG),
-            steps=int(smc_steps),
-            kl_weight=float(smc_kl_weight),
-            lambda_tempering=bool(smc_lambda_tempering),
-            lambda_one_at=float(smc_lambda_one_at),
-            num_batches=int(smc_num_batches),
-            num_particles=int(smc_num_particles),
-            proposal_type=str(smc_proposal_type),
-            use_continuous_formulation=bool(smc_use_continuous_formulation),
-            phi=int(smc_phi),
-            tau=float(smc_tau),
+            negative_prompt=smc_grad_negative_prompt or "",
+            ess_threshold=float(smc_grad_ess_threshold),
+            partial_resampling=bool(smc_grad_partial_resampling),
+            resample_frequency=int(smc_grad_resample_frequency),
+            CFG=float(smc_grad_CFG),
+            steps=int(smc_grad_steps),
+            kl_weight=float(smc_grad_kl_weight),
+            lambda_tempering=bool(smc_grad_lambda_tempering),
+            lambda_one_at=float(smc_grad_lambda_one_at),
+            num_particles=int(smc_grad_num_particles),
+            use_continuous_formulation=bool(smc_grad_use_continuous_formulation),
+            phi=int(smc_grad_phi),
+            tau=float(smc_grad_tau),
         )
-        smc_output = infer_smc_grad(smc_cfg, device=DEVICE)
-        # The above line is defensive; simpler: pass smc_device value used by gradio - will be provided.
+        smc_grad_output = infer_smc_grad(smc_grad_cfg, device=DEVICE)
+        # The above line is defensive; simpler: pass smc_grad_device value used by gradio - will be provided.
     except Exception as e:
-        smc_images = []
-        smc_output = None
-        smc_error = f"SMC inference error: {e}"
-        smc_images = [smc_error]
+        smc_grad_images = []
+        smc_grad_output = None
+        smc_grad_error = f"SMC inference error: {e}"
+        smc_grad_images = [smc_grad_error]
 
     # If outputs are dataclasses with PIL images, gr.Gallery accepts lists of PIL images.
     pretrained_gallery = pretrained_images if isinstance(pretrained_images, list) else [pretrained_images]
-    smc_gallery = smc_output.images if smc_output is not None else smc_images
+    smc_grad_gallery = smc_grad_output.images if smc_grad_output is not None else smc_grad_images
 
     pretrained_info = _format_inference_output(pretrained_output)
-    smc_info = _format_inference_output(smc_output)
+    smc_grad_info = _format_inference_output(smc_grad_output)
 
-    return pretrained_gallery, pretrained_info, smc_gallery, smc_info
+    return pretrained_gallery, pretrained_info, smc_grad_gallery, smc_grad_info
 
 
 with gr.Blocks() as demo:
@@ -174,42 +137,35 @@ with gr.Blocks() as demo:
     with gr.Row():
         with gr.Column(scale=1, min_width=280):
             with gr.Accordion("Pretrained method — settings", open=False):
-                pretrained_negative_prompt = gr.Textbox(label="Negative prompt", value="", lines=1)
-                pretrained_resolution = gr.Slider(256, MAX_IMAGE_SIZE, step=64, value=DEFAULTS["resolution"], label="Resolution")
-                pretrained_CFG = gr.Slider(0.0, 30.0, step=0.1, value=DEFAULTS["pretrained_CFG"], label="CFG")
-                pretrained_steps = gr.Slider(1, 200, step=1, value=DEFAULTS["pretrained_steps"], label="Steps")
-                pretrained_num_batches = gr.Slider(1, 8, step=1, value=DEFAULTS["pretrained_num_batches"], label="Batches")
-                pretrained_device = gr.Dropdown(choices=["cpu", "cuda"], value=("cuda" if torch.cuda.is_available() else "cpu"), label="Device")
+                pretrained_negative_prompt = gr.Textbox(label="Negative prompt", value=PretrainedInferenceConfig.negative_prompt, lines=1)
+                pretrained_CFG = gr.Slider(0.0, 30.0, step=0.1, value=PretrainedInferenceConfig.CFG, label="CFG")
+                pretrained_steps = gr.Slider(1, 200, step=1, value=PretrainedInferenceConfig.steps, label="Steps")
 
         with gr.Column(scale=2):
-            pretrained_gallery = gr.Gallery(label="Pretrained outputs", show_label=True, elem_id="pretrained_gallery", height="auto", columns=4)
+            pretrained_gallery = gr.Gallery(label="Pretrained outputs", show_label=True, elem_id="pretrained_gallery", height="70vw", columns=4)
             pretrained_info = gr.Textbox(label="Pretrained info", interactive=False)
 
     # --- SMC-grad method row ---
     with gr.Row():
         with gr.Column(scale=1, min_width=280):
             with gr.Accordion("SMC-grad method — settings", open=False):
-                smc_negative_prompt = gr.Textbox(label="Negative prompt", value="", lines=1)
-                smc_resolution = gr.Slider(256, MAX_IMAGE_SIZE, step=64, value=DEFAULTS["resolution"], label="Resolution")
-                smc_CFG = gr.Slider(0.0, 30.0, step=0.1, value=DEFAULTS["smc_CFG"], label="CFG")
-                smc_steps = gr.Slider(1, 200, step=1, value=DEFAULTS["smc_steps"], label="Steps")
-                smc_num_batches = gr.Slider(1, 8, step=1, value=DEFAULTS["smc_num_batches"], label="Batches")
-                smc_num_particles = gr.Slider(1, 64, step=1, value=DEFAULTS["smc_num_particles"], label="Num particles")
-                smc_ess_threshold = gr.Slider(0.0, 1.0, step=0.01, value=DEFAULTS["smc_ess_threshold"], label="ESS threshold")
-                smc_partial_resampling = gr.Checkbox(label="Partial resampling", value=DEFAULTS["smc_partial_resampling"])
-                smc_resample_frequency = gr.Slider(1, 50, step=1, value=DEFAULTS["smc_resample_frequency"], label="Resample frequency")
-                smc_kl_weight = gr.Slider(0.0, 10.0, step=0.01, value=DEFAULTS["smc_kl_weight"], label="KL weight")
-                smc_lambda_tempering = gr.Checkbox(label="Lambda tempering", value=DEFAULTS["smc_lambda_tempering"])
-                smc_lambda_one_at = gr.Slider(0.0, 1.0, step=0.01, value=DEFAULTS["smc_lambda_one_at"], label="Lambda one at (fraction of steps)")
-                smc_use_continuous_formulation = gr.Checkbox(label="Use continuous formulation", value=True)
-                smc_phi = gr.Slider(1, 8, step=1, value=DEFAULTS["smc_phi"], label="Phi")
-                smc_tau = gr.Slider(0.0, 1.0, step=0.001, value=DEFAULTS["smc_tau"], label="Tau")
-                smc_proposal_type = gr.Dropdown(choices=["locally_optimal", "without_SMC", "other"], value="locally_optimal", label="Proposal type")
-                smc_device = gr.Dropdown(choices=["cpu", "cuda"], value=("cuda" if torch.cuda.is_available() else "cpu"), label="Device")
+                smc_grad_negative_prompt = gr.Textbox(label="Negative prompt", value=SMCGradInferenceConfig.negative_prompt, lines=1)
+                smc_grad_CFG = gr.Slider(0.0, 30.0, step=0.1, value=SMCGradInferenceConfig.CFG, label="CFG")
+                smc_grad_steps = gr.Slider(1, 200, step=1, value=SMCGradInferenceConfig.steps, label="Steps")
+                smc_grad_num_particles = gr.Slider(1, 64, step=1, value=SMCGradInferenceConfig.num_particles, label="SMC Num particles")
+                smc_grad_ess_threshold = gr.Slider(0.0, 1.0, step=0.01, value=SMCGradInferenceConfig.ess_threshold, label="ESS threshold")
+                smc_grad_partial_resampling = gr.Checkbox(label="Partial resampling", value=SMCGradInferenceConfig.partial_resampling)
+                smc_grad_resample_frequency = gr.Slider(1, 50, step=1, value=SMCGradInferenceConfig.resample_frequency, label="Resample frequency")
+                smc_grad_kl_weight = gr.Slider(0.0, 10.0, step=0.01, value=SMCGradInferenceConfig.kl_weight, label="KL weight")
+                smc_grad_lambda_tempering = gr.Checkbox(label="Lambda tempering", value=SMCGradInferenceConfig.lambda_tempering)
+                smc_grad_lambda_one_at = gr.Slider(0.0, 1.0, step=0.01, value=SMCGradInferenceConfig.lambda_one_at, label="Lambda one at (fraction of steps)")
+                smc_grad_use_continuous_formulation = gr.Checkbox(label="Use continuous formulation", value=SMCGradInferenceConfig.use_continuous_formulation)
+                smc_grad_phi = gr.Slider(1, 8, step=1, value=SMCGradInferenceConfig.phi, label="Phi")
+                smc_grad_tau = gr.Slider(0.0, 1.0, step=0.001, value=SMCGradInferenceConfig.tau, label="Tau")
 
         with gr.Column(scale=2):
-            smc_gallery = gr.Gallery(label="SMC-grad outputs", show_label=True, elem_id="smc_gallery", height="auto", columns=4)
-            smc_info = gr.Textbox(label="SMC-grad info", interactive=False)
+            smc_grad_gallery = gr.Gallery(label="SMC-grad outputs", show_label=True, elem_id="smc_grad_gallery", height="70vw", columns=4)
+            smc_grad_info = gr.Textbox(label="SMC-grad info", interactive=False)
 
     # Wire up the run button and prompt submit to the same runner
     run_button.click(
@@ -218,30 +174,24 @@ with gr.Blocks() as demo:
             prompt,
 
             pretrained_negative_prompt,
-            pretrained_resolution,
             pretrained_CFG,
             pretrained_steps,
-            pretrained_num_batches,
-            pretrained_device,
-
-            smc_negative_prompt,
-            smc_resolution,
-            smc_CFG,
-            smc_steps,
-            smc_num_batches,
-            smc_num_particles,
-            smc_ess_threshold,
-            smc_partial_resampling,
-            smc_resample_frequency,
-            smc_kl_weight,
-            smc_lambda_tempering,
-            smc_lambda_one_at,
-            smc_use_continuous_formulation,
-            smc_phi,
-            smc_tau,
-            smc_proposal_type,
+
+            smc_grad_negative_prompt,
+            smc_grad_CFG,
+            smc_grad_steps,
+            smc_grad_num_particles,
+            smc_grad_ess_threshold,
+            smc_grad_partial_resampling,
+            smc_grad_resample_frequency,
+            smc_grad_kl_weight,
+            smc_grad_lambda_tempering,
+            smc_grad_lambda_one_at,
+            smc_grad_use_continuous_formulation,
+            smc_grad_phi,
+            smc_grad_tau,
         ],
-        outputs=[pretrained_gallery, pretrained_info, smc_gallery, smc_info],
+        outputs=[pretrained_gallery, pretrained_info, smc_grad_gallery, smc_grad_info],
     )
 
     # Also allow pressing Enter in the prompt to trigger
@@ -251,32 +201,26 @@ with gr.Blocks() as demo:
             prompt,
 
             pretrained_negative_prompt,
-            pretrained_resolution,
             pretrained_CFG,
             pretrained_steps,
-            pretrained_num_batches,
-            pretrained_device,
-
-            smc_negative_prompt,
-            smc_resolution,
-            smc_CFG,
-            smc_steps,
-            smc_num_batches,
-            smc_num_particles,
-            smc_ess_threshold,
-            smc_partial_resampling,
-            smc_resample_frequency,
-            smc_kl_weight,
-            smc_lambda_tempering,
-            smc_lambda_one_at,
-            smc_use_continuous_formulation,
-            smc_phi,
-            smc_tau,
-            smc_proposal_type,
+
+            smc_grad_negative_prompt,
+            smc_grad_CFG,
+            smc_grad_steps,
+            smc_grad_num_particles,
+            smc_grad_ess_threshold,
+            smc_grad_partial_resampling,
+            smc_grad_resample_frequency,
+            smc_grad_kl_weight,
+            smc_grad_lambda_tempering,
+            smc_grad_lambda_one_at,
+            smc_grad_use_continuous_formulation,
+            smc_grad_phi,
+            smc_grad_tau,
         ],
-        outputs=[pretrained_gallery, pretrained_info, smc_gallery, smc_info],
+        outputs=[pretrained_gallery, pretrained_info, smc_grad_gallery, smc_grad_info],
     )
 
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(share=True)