Init app.py

app.py

@@ -1,23 +1,282 @@
-import os
-import sys
-import time
-import json
-from dataclasses import dataclass, asdict
-from typing import Optional, Dict, Any, List
-import torch
-from PIL import Image
+"""
+Gradio app to compare multiple inference methods for Monetico model.
+
+This file wires your existing inference functions (infer_pretrained, infer_smc_grad)
+into a single UI with one shared prompt and per-method collapsed setting panels.
+
+Place this file at repository root (next to src/) and run:
+    python app.py
+
+Notes:
+- The code assumes your module that contains infer_pretrained and infer_smc_grad
+  is importable (e.g. package root with src/ on PYTHONPATH). Adjust imports if needed.
+- Defaults provided are reasonable starting points; tweak as you like.
+"""
+
 import gradio as gr
+import torch
+from typing import List
+
+# Import your inference functions and dataclasses
+# 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,
+}
+
+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:
+        return "No output"
+    try:
+        rewards = out.image_rewards
+        mem = out.gpu_mem_used
+        return f"Rewards: {rewards} | GPU mem (GB): {mem:.3f}"
+    except Exception:
+        return "Could not parse inference output"
+
+
+def run_inference_all(
+    prompt,
+
+    # 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,
+):
+    """Wrapper that runs both inference methods and returns UI-friendly outputs.
+
+    Returns:
+      pretrained_images, pretrained_info, smc_images, smc_info
+    """
+    # --- Pretrained ---
+    pretrained_output = None
+    pretrained_images = []
+    try:
+        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_images = pretrained_output.images
+    except Exception as e:
+        pretrained_images = []
+        pretrained_output = None
+        pretrained_error = f"Pretrained inference error: {e}"
+        pretrained_images = [pretrained_error]
+
+    # --- SMC-grad ---
+    smc_output = None
+    smc_images = []
+    try:
+        smc_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),
+        )
+        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.
+    except Exception as e:
+        smc_images = []
+        smc_output = None
+        smc_error = f"SMC inference error: {e}"
+        smc_images = [smc_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
+
+    pretrained_info = _format_inference_output(pretrained_output)
+    smc_info = _format_inference_output(smc_output)
+
+    return pretrained_gallery, pretrained_info, smc_gallery, smc_info
+
+
+with gr.Blocks() as demo:
+    gr.Markdown("# Monetico — Multi-method Inference Playground")
+
+    with gr.Row():
+        prompt = gr.Textbox(label="Prompt", placeholder="Enter prompt here", value=examples[0], lines=1)
+        run_button = gr.Button("Run", variant="primary")
+
+    gr.Examples(examples=examples, inputs=prompt)
+
+    # --- Pretrained method row ---
+    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")
+
+        with gr.Column(scale=2):
+            pretrained_gallery = gr.Gallery(label="Pretrained outputs", show_label=True, elem_id="pretrained_gallery", height="auto", 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")
+
+        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)
+
+    # Wire up the run button and prompt submit to the same runner
+    run_button.click(
+        fn=run_inference_all,
+        inputs=[
+            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,
+        ],
+        outputs=[pretrained_gallery, pretrained_info, smc_gallery, smc_info],
+    )
+
+    # Also allow pressing Enter in the prompt to trigger
+    prompt.submit(
+        fn=run_inference_all,
+        inputs=[
+            prompt,
+
+            pretrained_negative_prompt,
+            pretrained_resolution,
+            pretrained_CFG,
+            pretrained_steps,
+            pretrained_num_batches,
+            pretrained_device,
 
-from src.smc.transformer import Transformer2DModel
-from src.smc.pipeline import Pipeline
-from src.meissonic.scheduler import Scheduler
-from src.smc.scheduler import ReMDMScheduler, MeissonicScheduler
-from transformers import CLIPTextModelWithProjection, CLIPTokenizer
-from diffusers.models.autoencoders.vq_model import VQModel
-import src.smc.rewards as rewards
-from src.smc.resampling import resample
+            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,
+        ],
+        outputs=[pretrained_gallery, pretrained_info, smc_gallery, smc_info],
+    )
 
-device = "cuda"
-dtype = torch.bfloat16
-model_path = "Collov-Labs/Monetico"
 
+if __name__ == "__main__":
+    demo.launch()