vis

app.py

@@ -1,402 +1,249 @@
-# import gradio as gr
+import gradio as gr
 
-# from absl import flags
-# from absl import app
-# from ml_collections import config_flags
-# import os
+from absl import flags
+from absl import app
+from ml_collections import config_flags
+import os
 
-# import spaces #[uncomment to use ZeroGPU]
-# import torch
+import spaces #[uncomment to use ZeroGPU]
+import torch
 
 
-# import os
-# import random
+import os
+import random
 
-# import numpy as np
-# import torch
-# import torch.nn.functional as F
-# from torchvision.utils import save_image
-# from huggingface_hub import hf_hub_download
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torchvision.utils import save_image
+from huggingface_hub import hf_hub_download
 
-# from absl import logging
-# import ml_collections
+from absl import logging
+import ml_collections
 
-# from diffusion.flow_matching import ODEEulerFlowMatchingSolver
-# import utils
-# import libs.autoencoder
-# from libs.clip import FrozenCLIPEmbedder
-# from configs import t2i_512px_clip_dimr
+from diffusion.flow_matching import ODEEulerFlowMatchingSolver
+import utils
+import libs.autoencoder
+from libs.clip import FrozenCLIPEmbedder
+from configs import t2i_512px_clip_dimr
 
 
-# def unpreprocess(x: torch.Tensor) -> torch.Tensor:
-#     x = 0.5 * (x + 1.0)
-#     x.clamp_(0.0, 1.0)
-#     return x
+def unpreprocess(x: torch.Tensor) -> torch.Tensor:
+    x = 0.5 * (x + 1.0)
+    x.clamp_(0.0, 1.0)
+    return x
     
-# def cosine_similarity_torch(latent1: torch.Tensor, latent2: torch.Tensor) -> torch.Tensor:
-#     latent1_flat = latent1.view(-1)
-#     latent2_flat = latent2.view(-1)
-#     cosine_similarity = F.cosine_similarity(
-#         latent1_flat.unsqueeze(0), latent2_flat.unsqueeze(0), dim=1
-#     )
-#     return cosine_similarity
-
-# def kl_divergence(latent1: torch.Tensor, latent2: torch.Tensor) -> torch.Tensor:
-#     latent1_prob = F.softmax(latent1, dim=-1)
-#     latent2_prob = F.softmax(latent2, dim=-1)
-#     latent1_log_prob = torch.log(latent1_prob)
-#     kl_div = F.kl_div(latent1_log_prob, latent2_prob, reduction="batchmean")
-#     return kl_div
-
-# def batch_decode(_z: torch.Tensor, decode, batch_size: int = 10) -> torch.Tensor:
-#     num_samples = _z.size(0)
-#     decoded_batches = []
-
-#     for i in range(0, num_samples, batch_size):
-#         batch = _z[i : i + batch_size]
-#         decoded_batch = decode(batch)
-#         decoded_batches.append(decoded_batch)
-
-#     return torch.cat(decoded_batches, dim=0)
-
-# def get_caption(llm: str, text_model, prompt_dict: dict, batch_size: int):
-#     if batch_size == 3:
-#         # Only addition or only subtraction mode.
-#         assert len(prompt_dict) == 2, "Expected 2 prompts for batch_size 3."
-#         batch_prompts = list(prompt_dict.values()) + [" "]
-#     elif batch_size == 4:
-#         # Addition and subtraction mode.
-#         assert len(prompt_dict) == 3, "Expected 3 prompts for batch_size 4."
-#         batch_prompts = list(prompt_dict.values()) + [" "]
-#     elif batch_size >= 5:
-#         # Linear interpolation mode.
-#         assert len(prompt_dict) == 2, "Expected 2 prompts for linear interpolation."
-#         batch_prompts = [prompt_dict["prompt_1"]] + [" "] * (batch_size - 2) + [prompt_dict["prompt_2"]]
-#     else:
-#         raise ValueError(f"Unsupported batch_size: {batch_size}")
-
-#     if llm == "clip":
-#         latent, latent_and_others = text_model.encode(batch_prompts)
-#         context = latent_and_others["token_embedding"].detach()
-#     elif llm == "t5":
-#         latent, latent_and_others = text_model.get_text_embeddings(batch_prompts)
-#         context = (latent_and_others["token_embedding"] * 10.0).detach()
-#     else:
-#         raise NotImplementedError(f"Language model {llm} not supported.")
-
-#     token_mask = latent_and_others["token_mask"].detach()
-#     tokens = latent_and_others["tokens"].detach()
-#     captions = batch_prompts
-
-#     return context, token_mask, tokens, captions
-
-# # Load configuration and initialize models.
-# config_dict = t2i_512px_clip_dimr.get_config()
-# config = ml_collections.ConfigDict(config_dict)
-
-# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-# logging.info(f"Using device: {device}")
-
-# # Freeze configuration.
-# config = ml_collections.FrozenConfigDict(config)
-
-# torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
-# MAX_SEED = np.iinfo(np.int32).max
-# MAX_IMAGE_SIZE = 1024  # Currently not used.
-
-# # Load the main diffusion model.
-# repo_id = "QHL067/CrossFlow"
-# filename = "pretrained_models/t2i_512px_clip_dimr.pth"
-# checkpoint_path = hf_hub_download(repo_id=repo_id, filename=filename)
-# nnet = utils.get_nnet(**config.nnet)
-# nnet = nnet.to(device)
-# state_dict = torch.load(checkpoint_path, map_location=device)
-# nnet.load_state_dict(state_dict)
-# nnet.eval()
-
-# # Initialize text model.
-# llm = "clip"
-# clip = FrozenCLIPEmbedder()
-# clip.eval()
-# clip.to(device)
-
-# # Load autoencoder.
-# autoencoder = libs.autoencoder.get_model(**config.autoencoder)
-# autoencoder.to(device)
-
-
-# @torch.cuda.amp.autocast()
-# def encode(_batch: torch.Tensor) -> torch.Tensor:
-#     """Encode a batch of images using the autoencoder."""
-#     return autoencoder.encode(_batch)
-
-
-# @torch.cuda.amp.autocast()
-# def decode(_batch: torch.Tensor) -> torch.Tensor:
-#     """Decode a batch of latent vectors using the autoencoder."""
-#     return autoencoder.decode(_batch)
-
-
-# @spaces.GPU #[uncomment to use ZeroGPU]
-# def infer(
-#     prompt1,
-#     prompt2,
-#     seed,
-#     randomize_seed,
-#     guidance_scale,
-#     num_inference_steps,
-#     num_of_interpolation,
-#     save_gpu_memory=True,
-#     progress=gr.Progress(track_tqdm=True),
-# ):
-#     if randomize_seed:
-#         seed = random.randint(0, MAX_SEED)
-
-#     torch.manual_seed(seed)
-#     if device.type == "cuda":
-#         torch.cuda.manual_seed_all(seed)
-
-#     # Only support interpolation in this implementation.
-#     prompt_dict = {"prompt_1": prompt1, "prompt_2": prompt2}
-#     for key, value in prompt_dict.items():
-#         assert value is not None, f"{key} must not be None."
-#     assert num_of_interpolation >= 5, "For linear interpolation, please sample at least five images."
-
-#     # Get text embeddings and tokens.
-#     _context, _token_mask, _token, _caption = get_caption(
-#         llm, clip, prompt_dict=prompt_dict, batch_size=num_of_interpolation
-#     )
-
-#     with torch.no_grad():
-#         _z_gaussian = torch.randn(num_of_interpolation, *config.z_shape, device=device)
-#         _z_x0, _mu, _log_var = nnet(
-#             _context, text_encoder=True, shape=_z_gaussian.shape, mask=_token_mask
-#         )
-#         _z_init = _z_x0.reshape(_z_gaussian.shape)
-
-#         # Prepare the initial latent representations based on the number of interpolations.
-#         if num_of_interpolation == 3:
-#             # Addition or subtraction mode.
-#             if config.prompt_a is not None:
-#                 assert config.prompt_s is None, "Only one of prompt_a or prompt_s should be provided."
-#                 z_init_temp = _z_init[0] + _z_init[1]
-#             elif config.prompt_s is not None:
-#                 assert config.prompt_a is None, "Only one of prompt_a or prompt_s should be provided."
-#                 z_init_temp = _z_init[0] - _z_init[1]
-#             else:
-#                 raise NotImplementedError("Either prompt_a or prompt_s must be provided for 3-sample mode.")
-#             mean = z_init_temp.mean()
-#             std = z_init_temp.std()
-#             _z_init[2] = (z_init_temp - mean) / std
-
-#         elif num_of_interpolation == 4:
-#             z_init_temp = _z_init[0] + _z_init[1] - _z_init[2]
-#             mean = z_init_temp.mean()
-#             std = z_init_temp.std()
-#             _z_init[3] = (z_init_temp - mean) / std
-
-#         elif num_of_interpolation >= 5:
-#             tensor_a = _z_init[0]
-#             tensor_b = _z_init[-1]
-#             num_interpolations = num_of_interpolation - 2
-#             interpolations = [
-#                 tensor_a + (tensor_b - tensor_a) * (i / (num_interpolations + 1))
-#                 for i in range(1, num_interpolations + 1)
-#             ]
-#             _z_init = torch.stack([tensor_a] + interpolations + [tensor_b], dim=0)
-
-#         else:
-#             raise ValueError("Unsupported number of interpolations.")
-
-#         assert guidance_scale > 1, "Guidance scale must be greater than 1."
-
-#         has_null_indicator = hasattr(config.nnet.model_args, "cfg_indicator")
-#         ode_solver = ODEEulerFlowMatchingSolver(
-#             nnet,
-#             bdv_model_fn=None,
-#             step_size_type="step_in_dsigma",
-#             guidance_scale=guidance_scale,
-#         )
-#         _z, _ = ode_solver.sample(
-#             x_T=_z_init,
-#             batch_size=num_of_interpolation,
-#             sample_steps=num_inference_steps,
-#             unconditional_guidance_scale=guidance_scale,
-#             has_null_indicator=has_null_indicator,
-#         )
-
-#         if save_gpu_memory:
-#             image_unprocessed = batch_decode(_z, decode)
-#         else:
-#             image_unprocessed = decode(_z)
-
-#         samples = unpreprocess(image_unprocessed).contiguous()[0]
-
-#     # return samples, seed
-#     return seed
-
-
-# # examples = [
-# #     "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
-# #     "An astronaut riding a green horse",
-# #     "A delicious ceviche cheesecake slice",
-# # ]
+def cosine_similarity_torch(latent1: torch.Tensor, latent2: torch.Tensor) -> torch.Tensor:
+    latent1_flat = latent1.view(-1)
+    latent2_flat = latent2.view(-1)
+    cosine_similarity = F.cosine_similarity(
+        latent1_flat.unsqueeze(0), latent2_flat.unsqueeze(0), dim=1
+    )
+    return cosine_similarity
+
+def kl_divergence(latent1: torch.Tensor, latent2: torch.Tensor) -> torch.Tensor:
+    latent1_prob = F.softmax(latent1, dim=-1)
+    latent2_prob = F.softmax(latent2, dim=-1)
+    latent1_log_prob = torch.log(latent1_prob)
+    kl_div = F.kl_div(latent1_log_prob, latent2_prob, reduction="batchmean")
+    return kl_div
+
+def batch_decode(_z: torch.Tensor, decode, batch_size: int = 10) -> torch.Tensor:
+    num_samples = _z.size(0)
+    decoded_batches = []
+
+    for i in range(0, num_samples, batch_size):
+        batch = _z[i : i + batch_size]
+        decoded_batch = decode(batch)
+        decoded_batches.append(decoded_batch)
+
+    return torch.cat(decoded_batches, dim=0)
+
+def get_caption(llm: str, text_model, prompt_dict: dict, batch_size: int):
+    if batch_size == 3:
+        # Only addition or only subtraction mode.
+        assert len(prompt_dict) == 2, "Expected 2 prompts for batch_size 3."
+        batch_prompts = list(prompt_dict.values()) + [" "]
+    elif batch_size == 4:
+        # Addition and subtraction mode.
+        assert len(prompt_dict) == 3, "Expected 3 prompts for batch_size 4."
+        batch_prompts = list(prompt_dict.values()) + [" "]
+    elif batch_size >= 5:
+        # Linear interpolation mode.
+        assert len(prompt_dict) == 2, "Expected 2 prompts for linear interpolation."
+        batch_prompts = [prompt_dict["prompt_1"]] + [" "] * (batch_size - 2) + [prompt_dict["prompt_2"]]
+    else:
+        raise ValueError(f"Unsupported batch_size: {batch_size}")
+
+    if llm == "clip":
+        latent, latent_and_others = text_model.encode(batch_prompts)
+        context = latent_and_others["token_embedding"].detach()
+    elif llm == "t5":
+        latent, latent_and_others = text_model.get_text_embeddings(batch_prompts)
+        context = (latent_and_others["token_embedding"] * 10.0).detach()
+    else:
+        raise NotImplementedError(f"Language model {llm} not supported.")
+
+    token_mask = latent_and_others["token_mask"].detach()
+    tokens = latent_and_others["tokens"].detach()
+    captions = batch_prompts
+
+    return context, token_mask, tokens, captions
+
+# Load configuration and initialize models.
+config_dict = t2i_512px_clip_dimr.get_config()
+config = ml_collections.ConfigDict(config_dict)
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+logging.info(f"Using device: {device}")
+
+# Freeze configuration.
+config = ml_collections.FrozenConfigDict(config)
+
+torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+MAX_SEED = np.iinfo(np.int32).max
+MAX_IMAGE_SIZE = 1024  # Currently not used.
 
-# examples = [
-#     ["A dog cooking dinner in the kitchen", "An orange cat wearing sunglasses on a ship"],
-# ]
+# Load the main diffusion model.
+repo_id = "QHL067/CrossFlow"
+filename = "pretrained_models/t2i_512px_clip_dimr.pth"
+checkpoint_path = hf_hub_download(repo_id=repo_id, filename=filename)
+nnet = utils.get_nnet(**config.nnet)
+nnet = nnet.to(device)
+state_dict = torch.load(checkpoint_path, map_location=device)
+nnet.load_state_dict(state_dict)
+nnet.eval()
 
-# css = """
-# #col-container {
-#     margin: 0 auto;
-#     max-width: 640px;
-# }
-# """
-
-# with gr.Blocks(css=css) as demo:
-#     with gr.Column(elem_id="col-container"):
-#         gr.Markdown(" # CrossFlow")
-#         gr.Markdown(" CrossFlow directly transforms text representations into images for text-to-image generation, enabling interpolation in the input text latent space.")
-
-#         with gr.Row():
-#             prompt1 = gr.Text(
-#                 label="Prompt_1",
-#                 show_label=False,
-#                 max_lines=1,
-#                 placeholder="Enter your prompt for the first image",
-#                 container=False,
-#             )
-        
-#         with gr.Row():
-#             prompt2 = gr.Text(
-#                 label="Prompt_2",
-#                 show_label=False,
-#                 max_lines=1,
-#                 placeholder="Enter your prompt for the second image",
-#                 container=False,
-#             )
-
-#         with gr.Row():
-#             run_button = gr.Button("Run", scale=0, variant="primary")
-
-#         result = gr.Image(label="Result", show_label=False)
-
-#         with gr.Accordion("Advanced Settings", open=False):
-#             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=10.0,
-#                     step=0.1,
-#                     value=7.0,  # Replace with defaults that work for your model
-#                 )
-#             with gr.Row():
-#                 num_inference_steps = gr.Slider(
-#                     label="Number of inference steps",
-#                     minimum=1,
-#                     maximum=50,
-#                     step=1,
-#                     value=50,  # Replace with defaults that work for your model
-#                 )
-#             with gr.Row():
-#                 num_of_interpolation = gr.Slider(
-#                     label="Number of images for interpolation",
-#                     minimum=5,
-#                     maximum=50,
-#                     step=1,
-#                     value=10,  # Replace with defaults that work for your model
-#                 )
-
-#         gr.Examples(examples=examples, inputs=[prompt1, prompt2])
-#     gr.on(
-#         triggers=[run_button.click, prompt1.submit, prompt2.submit],
-#         fn=infer,
-#         inputs=[
-#             prompt1,
-#             prompt2,
-#             seed,
-#             randomize_seed,
-#             guidance_scale,
-#             num_inference_steps,
-#             num_of_interpolation,
-#         ],
-#         # outputs=[result, seed],
-#         outputs=[seed],
-#     )
-
-# if __name__ == "__main__":
-#     demo.launch()
+# Initialize text model.
+llm = "clip"
+clip = FrozenCLIPEmbedder()
+clip.eval()
+clip.to(device)
 
-import gradio as gr
-import numpy as np
-import random
+# Load autoencoder.
+autoencoder = libs.autoencoder.get_model(**config.autoencoder)
+autoencoder.to(device)
 
-# import spaces #[uncomment to use ZeroGPU]
-from diffusers import DiffusionPipeline
-import torch
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model_repo_id = "stabilityai/sdxl-turbo"  # Replace to the model you would like to use
+@torch.cuda.amp.autocast()
+def encode(_batch: torch.Tensor) -> torch.Tensor:
+    """Encode a batch of images using the autoencoder."""
+    return autoencoder.encode(_batch)
 
-if torch.cuda.is_available():
-    torch_dtype = torch.float16
-else:
-    torch_dtype = torch.float32
 
-pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
-pipe = pipe.to(device)
-
-MAX_SEED = np.iinfo(np.int32).max
-MAX_IMAGE_SIZE = 1024
+@torch.cuda.amp.autocast()
+def decode(_batch: torch.Tensor) -> torch.Tensor:
+    """Decode a batch of latent vectors using the autoencoder."""
+    return autoencoder.decode(_batch)
 
 
-# @spaces.GPU #[uncomment to use ZeroGPU]
+@spaces.GPU #[uncomment to use ZeroGPU]
 def infer(
-    prompt,
-    negative_prompt,
+    prompt1,
+    prompt2,
     seed,
     randomize_seed,
-    width,
-    height,
     guidance_scale,
     num_inference_steps,
+    num_of_interpolation,
+    save_gpu_memory=True,
     progress=gr.Progress(track_tqdm=True),
 ):
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
 
-    generator = torch.Generator().manual_seed(seed)
+    torch.manual_seed(seed)
+    if device.type == "cuda":
+        torch.cuda.manual_seed_all(seed)
 
-    image = pipe(
-        prompt=prompt,
-        negative_prompt=negative_prompt,
-        guidance_scale=guidance_scale,
-        num_inference_steps=num_inference_steps,
-        width=width,
-        height=height,
-        generator=generator,
-    ).images[0]
+    # Only support interpolation in this implementation.
+    prompt_dict = {"prompt_1": prompt1, "prompt_2": prompt2}
+    for key, value in prompt_dict.items():
+        assert value is not None, f"{key} must not be None."
+    assert num_of_interpolation >= 5, "For linear interpolation, please sample at least five images."
 
-    print('image.shape')
-    print(image.shape)
+    # Get text embeddings and tokens.
+    _context, _token_mask, _token, _caption = get_caption(
+        llm, clip, prompt_dict=prompt_dict, batch_size=num_of_interpolation
+    )
 
-    return image, seed
+    with torch.no_grad():
+        _z_gaussian = torch.randn(num_of_interpolation, *config.z_shape, device=device)
+        _z_x0, _mu, _log_var = nnet(
+            _context, text_encoder=True, shape=_z_gaussian.shape, mask=_token_mask
+        )
+        _z_init = _z_x0.reshape(_z_gaussian.shape)
+
+        # Prepare the initial latent representations based on the number of interpolations.
+        if num_of_interpolation == 3:
+            # Addition or subtraction mode.
+            if config.prompt_a is not None:
+                assert config.prompt_s is None, "Only one of prompt_a or prompt_s should be provided."
+                z_init_temp = _z_init[0] + _z_init[1]
+            elif config.prompt_s is not None:
+                assert config.prompt_a is None, "Only one of prompt_a or prompt_s should be provided."
+                z_init_temp = _z_init[0] - _z_init[1]
+            else:
+                raise NotImplementedError("Either prompt_a or prompt_s must be provided for 3-sample mode.")
+            mean = z_init_temp.mean()
+            std = z_init_temp.std()
+            _z_init[2] = (z_init_temp - mean) / std
+
+        elif num_of_interpolation == 4:
+            z_init_temp = _z_init[0] + _z_init[1] - _z_init[2]
+            mean = z_init_temp.mean()
+            std = z_init_temp.std()
+            _z_init[3] = (z_init_temp - mean) / std
+
+        elif num_of_interpolation >= 5:
+            tensor_a = _z_init[0]
+            tensor_b = _z_init[-1]
+            num_interpolations = num_of_interpolation - 2
+            interpolations = [
+                tensor_a + (tensor_b - tensor_a) * (i / (num_interpolations + 1))
+                for i in range(1, num_interpolations + 1)
+            ]
+            _z_init = torch.stack([tensor_a] + interpolations + [tensor_b], dim=0)
+
+        else:
+            raise ValueError("Unsupported number of interpolations.")
+
+        assert guidance_scale > 1, "Guidance scale must be greater than 1."
+
+        has_null_indicator = hasattr(config.nnet.model_args, "cfg_indicator")
+        ode_solver = ODEEulerFlowMatchingSolver(
+            nnet,
+            bdv_model_fn=None,
+            step_size_type="step_in_dsigma",
+            guidance_scale=guidance_scale,
+        )
+        _z, _ = ode_solver.sample(
+            x_T=_z_init,
+            batch_size=num_of_interpolation,
+            sample_steps=num_inference_steps,
+            unconditional_guidance_scale=guidance_scale,
+            has_null_indicator=has_null_indicator,
+        )
+
+        if save_gpu_memory:
+            image_unprocessed = batch_decode(_z, decode)
+        else:
+            image_unprocessed = decode(_z)
+
+        samples = unpreprocess(image_unprocessed).contiguous()[0]
+
+    # return samples, seed
+    return seed
 
 
+# examples = [
+#     "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
+#     "An astronaut riding a green horse",
+#     "A delicious ceviche cheesecake slice",
+# ]
+
 examples = [
-    "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
-    "An astronaut riding a green horse",
-    "A delicious ceviche cheesecake slice",
+    ["A dog cooking dinner in the kitchen", "An orange cat wearing sunglasses on a ship"],
 ]
 
 css = """
@@ -408,29 +255,33 @@ css = """
 
 with gr.Blocks(css=css) as demo:
     with gr.Column(elem_id="col-container"):
-        gr.Markdown(" # Text-to-Image Gradio Template")
+        gr.Markdown(" # CrossFlow")
+        gr.Markdown(" CrossFlow directly transforms text representations into images for text-to-image generation, enabling interpolation in the input text latent space.")
 
         with gr.Row():
-            prompt = gr.Text(
-                label="Prompt",
+            prompt1 = gr.Text(
+                label="Prompt_1",
                 show_label=False,
                 max_lines=1,
-                placeholder="Enter your prompt",
+                placeholder="Enter your prompt for the first image",
+                container=False,
+            )
+        
+        with gr.Row():
+            prompt2 = gr.Text(
+                label="Prompt_2",
+                show_label=False,
+                max_lines=1,
+                placeholder="Enter your prompt for the second image",
                 container=False,
             )
 
+        with gr.Row():
             run_button = gr.Button("Run", scale=0, variant="primary")
 
         result = gr.Image(label="Result", show_label=False)
 
         with gr.Accordion("Advanced Settings", open=False):
-            negative_prompt = gr.Text(
-                label="Negative prompt",
-                max_lines=1,
-                placeholder="Enter a negative prompt",
-                visible=False,
-            )
-
             seed = gr.Slider(
                 label="Seed",
                 minimum=0,
@@ -441,56 +292,47 @@ with gr.Blocks(css=css) as demo:
 
             randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
 
-            with gr.Row():
-                width = gr.Slider(
-                    label="Width",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=1024,  # Replace with defaults that work for your model
-                )
-
-                height = gr.Slider(
-                    label="Height",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=1024,  # Replace with defaults that work for your model
-                )
-
             with gr.Row():
                 guidance_scale = gr.Slider(
                     label="Guidance scale",
                     minimum=0.0,
                     maximum=10.0,
                     step=0.1,
-                    value=0.0,  # Replace with defaults that work for your model
+                    value=7.0,  # Replace with defaults that work for your model
                 )
-
+            with gr.Row():
                 num_inference_steps = gr.Slider(
                     label="Number of inference steps",
                     minimum=1,
                     maximum=50,
                     step=1,
-                    value=2,  # Replace with defaults that work for your model
+                    value=50,  # Replace with defaults that work for your model
+                )
+            with gr.Row():
+                num_of_interpolation = gr.Slider(
+                    label="Number of images for interpolation",
+                    minimum=5,
+                    maximum=50,
+                    step=1,
+                    value=10,  # Replace with defaults that work for your model
                 )
 
-        gr.Examples(examples=examples, inputs=[prompt])
+        gr.Examples(examples=examples, inputs=[prompt1, prompt2])
     gr.on(
-        triggers=[run_button.click, prompt.submit],
+        triggers=[run_button.click, prompt1.submit, prompt2.submit],
         fn=infer,
         inputs=[
-            prompt,
-            negative_prompt,
+            prompt1,
+            prompt2,
             seed,
             randomize_seed,
-            width,
-            height,
             guidance_scale,
             num_inference_steps,
+            num_of_interpolation,
         ],
-        outputs=[result, seed],
+        # outputs=[result, seed],
+        outputs=[seed],
     )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()