More fixes for demo.

apps/shark_studio/api/initializers.py

@@ -25,6 +25,14 @@ def imports():
     )
     warnings.filterwarnings(action="ignore", category=UserWarning, module="torchvision")
     warnings.filterwarnings(action="ignore", category=UserWarning, module="torch")
+    warnings.filterwarnings(action="ignore", category=UserWarning, module="diffusers")
+    warnings.filterwarnings(action="ignore", category=FutureWarning, module="diffusers")
+    warnings.filterwarnings(
+        action="ignore", category=FutureWarning, module="huggingface-hub"
+    )
+    warnings.filterwarnings(
+        action="ignore", category=UserWarning, module="huggingface-hub"
+    )
 
     import gradio  # noqa: F401
 

apps/shark_studio/api/sd.py

@@ -99,7 +99,10 @@ class StableDiffusion:
         import_ir: bool = True,
         is_controlled: bool = False,
         external_weights: str = "safetensors",
+        progress=gr.Progress(),
     ):
+        progress(None, desc="Initializing pipeline...")
+        self.ui_device = device
         self.precision = precision
         self.compiled_pipeline = False
         self.base_model_id = base_model_id
@@ -178,13 +181,19 @@ class StableDiffusion:
             external_weights=external_weights,
             custom_vae=custom_vae,
         )
-        print(f"\n[LOG] Pipeline initialized with pipe_id: {self.pipe_id}.")
+        progress(None, desc="Pipeline initialized!...")
         gc.collect()
 
     def prepare_pipe(
-        self, custom_weights, adapters, embeddings, is_img2img, compiled_pipeline
+        self,
+        custom_weights,
+        adapters,
+        embeddings,
+        is_img2img,
+        compiled_pipeline,
+        progress=gr.Progress(),
     ):
-        print(f"\n[LOG] Preparing pipeline...")
+
         self.is_img2img = False
         mlirs = copy.deepcopy(self.model_map)
         vmfbs = copy.deepcopy(self.model_map)
@@ -236,17 +245,18 @@ class StableDiffusion:
                         "diffusion_pytorch_model.safetensors",
                     )
                     weights[key] = save_irpa(vae_weights_path, "vae.")
+        progress(None, desc=f"Preparing pipeline for {self.ui_device}...")
 
         vmfbs, weights = self.sd_pipe.check_prepared(
             mlirs, vmfbs, weights, interactive=False
         )
-        print(f"\n[LOG] Loading pipeline to device {self.rt_device}.")
+        progress(None, desc=f"Artifacts ready!")
+        progress(None, desc=f"Loading pipeline on device {self.ui_device}...")
+
         self.sd_pipe.load_pipeline(
             vmfbs, weights, self.rt_device, self.compiled_pipeline
         )
-        print(
-            "\n[LOG] Pipeline successfully prepared for runtime. Generating images..."
-        )
+        progress(None, desc="Pipeline loaded! Generating images...")
         return
 
     def generate_images(
@@ -261,7 +271,9 @@ class StableDiffusion:
         resample_type,
         control_mode,
         hints,
+        progress=gr.Progress()
     ):
+
         img = self.sd_pipe.generate_images(
             prompt,
             negative_prompt,
@@ -273,9 +285,7 @@ class StableDiffusion:
         return img
 
 
-def shark_sd_fn_dict_input(
-    sd_kwargs: dict,
-):
+def shark_sd_fn_dict_input(sd_kwargs: dict, *, progress=gr.Progress()):
     print("\n[LOG] Submitting Request...")
 
     for key in sd_kwargs:
@@ -340,6 +350,8 @@ def shark_sd_fn(
     resample_type: str,
     controlnets: dict,
     embeddings: dict,
+    seed_increment: str | int = 1,
+    progress=gr.Progress(),
 ):
     sd_kwargs = locals()
     if not isinstance(sd_init_image, list):
@@ -438,6 +450,10 @@ def shark_sd_fn(
         global_obj.get_sd_obj().prepare_pipe(**submit_prep_kwargs)
 
     generated_imgs = []
+    if seed == -1:
+        seed = randint(0, sys.maxsize)
+    progress(None, desc=f"Generating...")
+
     for current_batch in range(batch_count):
         start_time = time.time()
         out_imgs = global_obj.get_sd_obj().generate_images(**submit_run_kwargs)
@@ -456,14 +472,20 @@ def shark_sd_fn(
                 sd_kwargs,
             )
         generated_imgs.extend(out_imgs)
-        # TODO: make seed changes over batch counts more configurable.
-        submit_run_kwargs["seed"] = submit_run_kwargs["seed"] + 1
+        seed = get_next_seed(seed, seed_increment)
         yield generated_imgs, status_label(
             "Stable Diffusion", current_batch + 1, batch_count, batch_size
         )
     return (generated_imgs, "")
 
 
+def get_next_seed(seed, seed_increment):
+    if isinstance(seed_increment, int):
+        return int(seed + seed_increment)
+    elif seed_increment == "random":
+        return randint(0, sys.maxsize)
+
+
 def unload_sd():
     print("Unloading models.")
     import apps.shark_studio.web.utils.globals as global_obj

apps/shark_studio/api/utils.py

@@ -54,31 +54,31 @@ def get_available_devices():
     available_devices = []
     rocm_devices = get_devices_by_name("rocm")
     available_devices.extend(rocm_devices)
-    cpu_device = get_devices_by_name("cpu-sync")
-    available_devices.extend(cpu_device)
+    # cpu_device = get_devices_by_name("cpu-sync")
+    # available_devices.extend(cpu_device)
     cpu_device = get_devices_by_name("cpu-task")
     available_devices.extend(cpu_device)
 
-    from shark.iree_utils.vulkan_utils import (
-        get_all_vulkan_devices,
-    )
+    # from shark.iree_utils.vulkan_utils import (
+    #     get_all_vulkan_devices,
+    # )
 
-    vulkaninfo_list = get_all_vulkan_devices()
-    vulkan_devices = []
-    id = 0
-    for device in vulkaninfo_list:
-        vulkan_devices.append(f"{device.strip()} => vulkan://{id}")
-        id += 1
-    if id != 0:
-        print(f"vulkan devices are available.")
+    # vulkaninfo_list = get_all_vulkan_devices()
+    # vulkan_devices = []
+    # id = 0
+    # for device in vulkaninfo_list:
+    #     vulkan_devices.append(f"{device.strip()} => vulkan://{id}")
+    #     id += 1
+    # if id != 0:
+    #     print(f"vulkan devices are available.")
 
-    available_devices.extend(vulkan_devices)
-    metal_devices = get_devices_by_name("metal")
-    available_devices.extend(metal_devices)
-    cuda_devices = get_devices_by_name("cuda")
-    available_devices.extend(cuda_devices)
-    hip_devices = get_devices_by_name("hip")
-    available_devices.extend(hip_devices)
+    # available_devices.extend(vulkan_devices)
+    # metal_devices = get_devices_by_name("metal")
+    # available_devices.extend(metal_devices)
+    # cuda_devices = get_devices_by_name("cuda")
+    # available_devices.extend(cuda_devices)
+    # hip_devices = get_devices_by_name("hip")
+    # available_devices.extend(hip_devices)
 
     for idx, device_str in enumerate(available_devices):
         if "AMD Radeon(TM) Graphics =>" in device_str:

apps/shark_studio/modules/ckpt_processing.py

@@ -71,7 +71,14 @@ def save_irpa(weights_path, prepend_str):
         new_key = prepend_str + key
         archive.add_tensor(new_key, weights[key])
 
-    irpa_file = weights_path.replace(".safetensors", ".irpa")
+    if "safetensors" in weights_path:
+        irpa_file = weights_path.replace(".safetensors", ".irpa")
+    elif "irpa" in weights_path:
+        irpa_file = weights_path
+    else:
+        return Exception(
+            "Invalid file format. Please provide a .safetensors or .irpa file."
+        )
     archive.save(irpa_file)
     return irpa_file
 

apps/shark_studio/modules/img_processing.py

@@ -33,6 +33,8 @@ def save_output_img(output_img, img_seed, extra_info=None):
 
     if extra_info is None:
         extra_info = {}
+    elif "progress" in extra_info.keys():
+        extra_info.pop("progress")
     generated_imgs_path = Path(
         get_generated_imgs_path(), get_generated_imgs_todays_subdir()
     )

apps/shark_studio/modules/shared_cmd_opts.py

@@ -35,10 +35,7 @@ p.add_argument(
     "--prompt",
     nargs="+",
     default=[
-        "a photo taken of the front of a super-car drifting on a road near "
-        "mountains at high speeds with smoke coming off the tires, front "
-        "angle, front point of view, trees in the mountains of the "
-        "background, ((sharp focus))"
+        "A hi-res photo of a red street racer drifting around a curve on a mountain, high altitude, at night, tokyo in the background, 8k"
     ],
     help="Text of which images to be generated.",
 )
@@ -62,7 +59,7 @@ p.add_argument(
 p.add_argument(
     "--steps",
     type=int,
-    default=50,
+    default=2,
     help="The number of steps to do the sampling.",
 )
 
@@ -100,7 +97,7 @@ p.add_argument(
 p.add_argument(
     "--guidance_scale",
     type=float,
-    default=7.5,
+    default=0,
     help="The value to be used for guidance scaling.",
 )
 
@@ -346,7 +343,7 @@ p.add_argument(
 p.add_argument(
     "--batch_count",
     type=int,
-    default=1,
+    default=4,
     help="Number of batches to be generated with random seeds in " "single execution.",
 )
 

apps/shark_studio/tools/params_prefixer.py

@@ -0,0 +1,20 @@
+from apps.shark_studio.modules.ckpt_processing import save_irpa
+import argparse
+import safetensors
+
+parser = argparse.ArgumentParser()
+parser.add_argument(
+    "--input",
+    type=str,
+    default="",
+    help="input safetensors/irpa",
+)
+parser.add_argument(
+    "--prefix",
+    type=str,
+    default="",
+    help="prefix to add to all the keys in the irpa",
+)
+args = parser.parse_args()
+output_file = save_irpa(args.input, args.prefix)
+print("saved irpa to", output_file, "with prefix", args.prefix)

apps/shark_studio/web/index.py

@@ -194,7 +194,7 @@ def webui():
                 sd_element.render()
             with gr.TabItem(label="Output Gallery", id=1):
                 outputgallery_element.render()
-            with gr.TabItem(label="Chat Bot", id=2):
+            with gr.TabItem(label="Chat Bot", id=2, visible=False):
                 chat_element.render()
 
     studio_web.queue()

apps/shark_studio/web/ui/sd.py

@@ -45,12 +45,13 @@ from apps.shark_studio.modules import logger
 import apps.shark_studio.web.utils.globals as global_obj
 
 sd_default_models = [
-    "runwayml/stable-diffusion-v1-5",
-    "stabilityai/stable-diffusion-2-1-base",
-    "stabilityai/stable-diffusion-2-1",
+    # "runwayml/stable-diffusion-v1-5",
+    # "stabilityai/stable-diffusion-2-1-base",
+    # "stabilityai/stable-diffusion-2-1",
     "stabilityai/stable-diffusion-xl-base-1.0",
     "stabilityai/sdxl-turbo",
 ]
+sd_default_models.extend(get_checkpoints(model_type="scripts"))
 
 
 def view_json_file(file_path):
@@ -158,6 +159,8 @@ def load_sd_cfg(sd_json: dict, load_sd_config: str):
                 sd_image = [Image.open(i, mode="r")]
     else:
         sd_image = None
+    if not sd_json["device"]:
+        sd_json["device"] = gr.update()
 
     return [
         sd_json["prompt"][0],
@@ -198,7 +201,7 @@ def save_sd_cfg(config: dict, save_name: str):
         filepath += ".json"
     with open(filepath, mode="w") as f:
         f.write(json.dumps(config))
-    return "..."
+    return save_name
 
 
 def create_canvas(width, height):
@@ -256,16 +259,17 @@ with gr.Blocks(title="Stable Diffusion") as sd_element:
                         choices=global_obj.get_device_list(),
                         allow_custom_value=False,
                     )
-                    target_triple = gr.Textbox(
-                        elem_id="target_triple",
-                        label="Architecture",
-                        value="",
-                    )
                     with gr.Row():
                         ondemand = gr.Checkbox(
                             value=cmd_opts.lowvram,
                             label="Low VRAM",
                             interactive=True,
+                            visible=False,
+                        )
+                        target_triple = gr.Textbox(
+                            elem_id="target_triple",
+                            label="Architecture",
+                            value="",
                         )
                         precision = gr.Radio(
                             label="Precision",
@@ -281,7 +285,7 @@ with gr.Blocks(title="Stable Diffusion") as sd_element:
                     label="\U000026F0\U0000FE0F Base Model",
                     info="Select or enter HF model ID",
                     elem_id="custom_model",
-                    value="stabilityai/stable-diffusion-2-1-base",
+                    value="stabilityai/sdxl-turbo",
                     choices=sd_default_models,
                     allow_custom_value=True,
                 )  # base_model_id
@@ -301,7 +305,9 @@ with gr.Blocks(title="Stable Diffusion") as sd_element:
                         label="\U00002194\U0000FE0F Width",
                     )
                 with gr.Accordion(
-                    label="\U00002696\U0000FE0F Model Weights", open=False
+                    label="\U00002696\U0000FE0F Model Weights",
+                    open=False,
+                    visible=False,  # DEMO
                 ):
                     with gr.Column():
                         custom_weights = gr.Dropdown(
@@ -369,7 +375,9 @@ with gr.Blocks(title="Stable Diffusion") as sd_element:
                             show_progress=False,
                         )
                 with gr.Accordion(
-                    label="\U0001F9EA\U0000FE0F Input Image Processing", open=False
+                    label="\U0001F9EA\U0000FE0F Input Image Processing",
+                    open=False,
+                    visible=False,
                 ):
                     strength = gr.Slider(
                         0,
@@ -403,21 +411,21 @@ with gr.Blocks(title="Stable Diffusion") as sd_element:
                     seed = gr.Textbox(
                         value=cmd_opts.seed,
                         label="\U0001F331\U0000FE0F Seed",
-                        info="An integer or a JSON list of integers, -1 for random",
+                        info="An integer, -1 for random",
                         show_copy_button=True,
                     )
                     scheduler = gr.Dropdown(
                         elem_id="scheduler",
                         label="\U0001F4C5\U0000FE0F Scheduler",
                         info="\U000E0020",  # forces same height as seed
-                        value="EulerDiscrete",
+                        value="EulerAncestralDiscrete",
                         choices=scheduler_model_map.keys(),
                         allow_custom_value=False,
                     )
                 with gr.Row():
                     steps = gr.Slider(
                         1,
-                        100,
+                        50,
                         value=cmd_opts.steps,
                         step=1,
                         label="\U0001F3C3\U0000FE0F Steps",
@@ -478,17 +486,17 @@ with gr.Blocks(title="Stable Diffusion") as sd_element:
                             with gr.Row():
                                 canvas_width = gr.Slider(
                                     label="Canvas Width",
-                                    minimum=256,
+                                    minimum=512,
                                     maximum=1024,
                                     value=512,
-                                    step=8,
+                                    step=512,
                                 )
                                 canvas_height = gr.Slider(
                                     label="Canvas Height",
-                                    minimum=256,
+                                    minimum=512,
                                     maximum=1024,
                                     value=512,
-                                    step=8,
+                                    step=512,
                                 )
                             make_canvas = gr.Button(
                                 value="Make Canvas!",
@@ -558,7 +566,9 @@ with gr.Blocks(title="Stable Diffusion") as sd_element:
                         lambda: gr.Tabs(selected=101),
                         outputs=[sd_tabs],
                     )
-                    with gr.Tab(label="Input Image", id=100) as sd_tab_init_image:
+                    with gr.Tab(
+                        label="Input Image", id=100, visible=False
+                    ) as sd_tab_init_image:  # DEMO
                         with gr.Column(elem_classes=["sd-right-panel"]):
                             with gr.Row(elem_classes=["fill"]):
                                 # TODO: make this import image prompt info if it exists
@@ -604,10 +614,10 @@ with gr.Blocks(title="Stable Diffusion") as sd_element:
                                     step=1,
                                     label="Batch Size",
                                     interactive=True,
-                                    visible=True,
+                                    visible=False,  # DEMO
                                 )
                                 compiled_pipeline = gr.Checkbox(
-                                    False,
+                                    True,
                                     label="Faster txt2img (SDXL only)",
                                 )
                             with gr.Row():
@@ -618,18 +628,18 @@ with gr.Blocks(title="Stable Diffusion") as sd_element:
                                     queue=False,
                                     show_progress=False,
                                 )
-                                stop_batch = gr.Button("Stop")
+                                stop_batch = gr.Button("Stop", visible=False)
                     with gr.Tab(label="Config", id=102) as sd_tab_config:
                         with gr.Column(elem_classes=["sd-right-panel"]):
                             with gr.Row(elem_classes=["fill"]):
                                 Path(get_configs_path()).mkdir(
                                     parents=True, exist_ok=True
                                 )
+                                write_default_sd_configs(get_configs_path())
                                 default_config_file = os.path.join(
                                     get_configs_path(),
-                                    "default_sd_config.json",
+                                    "sdxl-turbo.json",
                                 )
-                                write_default_sd_configs(get_configs_path())
                                 sd_json = gr.JSON(
                                     elem_classes=["fill"],
                                     value=view_json_file(default_config_file),
@@ -644,7 +654,7 @@ with gr.Blocks(title="Stable Diffusion") as sd_element:
                                             if cmd_opts.configs_path
                                             else get_configs_path()
                                         ),
-                                        height=75,
+                                        height=200,
                                     )
                                 with gr.Column(scale=1):
                                     save_sd_config = gr.Button(
@@ -655,13 +665,13 @@ with gr.Blocks(title="Stable Diffusion") as sd_element:
                                         size="sm",
                                         components=sd_json,
                                     )
-                            with gr.Row():
-                                sd_config_name = gr.Textbox(
-                                    value="Config Name",
-                                    info="Name of the file this config will be saved to.",
-                                    interactive=True,
-                                    show_label=False,
-                                )
+                            #with gr.Row():
+                                    sd_config_name = gr.Textbox(
+                                        value="Config Name",
+                                        info="Name of the file this config will be saved to.",
+                                        interactive=True,
+                                        show_label=False,
+                                    )
                                 load_sd_config.change(
                                     fn=load_sd_cfg,
                                     inputs=[sd_json, load_sd_config],
@@ -749,6 +759,7 @@ with gr.Blocks(title="Stable Diffusion") as sd_element:
         outputs=[
             sd_json,
         ],
+        show_progress=False,
     )
 
     status_kwargs = dict(

apps/shark_studio/web/utils/default_configs.py

@@ -89,7 +89,7 @@ sdxl_turbo = r"""{
 }"""
 
 default_sd_configs = {
-    "default_sd_config.json": default_sd_config,
+    # "default_sd_config.json": sdxl_turbo,
     "sdxl-30steps.json": sdxl_30steps,
     "sdxl-turbo.json": sdxl_turbo,
 }

apps/shark_studio/web/utils/file_utils.py

@@ -17,8 +17,9 @@ from apps.shark_studio.web.utils.default_configs import default_sd_configs
 def write_default_sd_configs(path):
     for key in default_sd_configs.keys():
         config_fpath = os.path.join(path, key)
-        with open(config_fpath, "w") as f:
-            f.write(default_sd_configs[key])
+        if not os.path.exists(config_fpath):
+            with open(config_fpath, "w") as f:
+                f.write(default_sd_configs[key])
 
 
 def safe_name(name):
@@ -87,6 +88,8 @@ def get_checkpoints_path(model_type=""):
 def get_checkpoints(model_type="checkpoints"):
     ckpt_files = []
     file_types = checkpoints_filetypes
+    if model_type == "scripts":
+        file_types = ["shark_*.py"]
     if model_type == "lora":
         file_types = file_types + ("*.pt", "*.bin")
     for extn in file_types:

setup_venv.ps1

@@ -89,7 +89,7 @@ else {python -m venv .\shark.venv\}
 python -m pip install --upgrade pip
 pip install wheel
 pip install --pre -r requirements.txt
-pip install --force-reinstall https://github.com/nod-ai/SRT/releases/download/candidate-20240528.279/iree_compiler-20240528.279-cp311-cp311-win_amd64.whl https://github.com/nod-ai/SRT/releases/download/candidate-20240528.279/iree_runtime-20240528.279-cp311-cp311-win_amd64.whl 
+pip install --force-reinstall https://github.com/nod-ai/SRT/releases/download/candidate-20240601.282/iree_compiler-20240601.282-cp311-cp311-win_amd64.whl https://github.com/nod-ai/SRT/releases/download/candidate-20240601.282/iree_runtime-20240601.282-cp311-cp311-win_amd64.whl
 pip install -e .
 
 Write-Host "Source your venv with ./shark.venv/Scripts/activate"