adding slider edits

app.py

@@ -8,6 +8,9 @@ import torch
 from diffusers import DiffusionPipeline, UNet2DConditionModel, LCMScheduler
 from huggingface_hub import hf_hub_download
 from safetensors.torch import load_file
+import sys
+sys.path.append('.')
+from utils.lora import LoRANetwork, DEFAULT_TARGET_REPLACE, UNET_TARGET_REPLACE_MODULE_CONV
 
 model_repo_id = "stabilityai/stable-diffusion-xl-base-1.0"
 repo_name = "tianweiy/DMD2"
@@ -27,6 +30,26 @@ pipe = DiffusionPipeline.from_pretrained(model_repo_id, unet=unet, torch_dtype=t
 pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
 
 
+unet = pipe.unet
+
+## Change these parameters based on how you trained your sliderspace sliders
+train_method = 'xattn-strict'
+rank = 1 
+alpha =1 
+networks = {}
+modules = DEFAULT_TARGET_REPLACE
+modules += UNET_TARGET_REPLACE_MODULE_CONV
+for i in range(1):
+    networks[i] = LoRANetwork(
+        unet,
+        rank=int(rank),
+        multiplier=1.0,
+        alpha=int(alpha),
+        train_method=train_method,
+        fast_init=True,
+    ).to(device, dtype=torch_dtype)
+
+
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
@@ -42,25 +65,49 @@ def infer(
     height,
     guidance_scale,
     num_inference_steps,
+    sliderspace_path,
+    slider_scale,
     progress=gr.Progress(track_tqdm=True),
 ):
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
 
-    generator = torch.Generator().manual_seed(seed)
+    for net in networks:
+        networks[net].load_state_dict(torch.load(sliderspace_path))
 
-    # with network:
+    for net in networks:
+        networks[net].set_lora_slider(slider_scale)
+
+    with networks[0]:
+        pass
+
+    # original image
+    generator = torch.Generator().manual_seed(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]
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            guidance_scale=guidance_scale,
+            num_inference_steps=num_inference_steps,
+            width=width,
+            height=height,
+            generator=generator,
+        ).images[0]
+
+    # edited image
+    generator = torch.Generator().manual_seed(seed)
+    with  networks[0]:
+        slider_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]
 
-    return image, seed
+    
+    return image, slider_image, seed
 
 
 examples = [
@@ -91,7 +138,32 @@ with gr.Blocks(css=css) as demo:
 
             run_button = gr.Button("Run", scale=0, variant="primary")
 
-        result = gr.Image(label="Result", show_label=False)
+
+        # New dropdowns side by side
+        with gr.Row():
+            slider_space = gr.Dropdown(
+                choices=["spaceship", "car", "person"],
+                label="SliderSpace",
+                value="spaceship"
+            )
+            discovered_directions = gr.Dropdown(
+                choices=[f"Slider {i}" for i in range(1, 11)],
+                label="Discovered Directions",
+                value="Slider 1"
+            )
+
+            slider_scale =  gr.Slider(
+                    label="Slider Scale",
+                    minimum=-2,
+                    maximum=2,
+                    step=0.1,
+                    value=1,  
+                )
+        
+        with gr.Row():
+            result = gr.Image(label="Original Image", show_label=True)
+            slider_result = gr.Image(label="Discovered Edit Direction", show_label=True)
+        
 
         with gr.Accordion("Advanced Settings", open=False):
             negative_prompt = gr.Text(
@@ -158,8 +230,10 @@ with gr.Blocks(css=css) as demo:
             height,
             guidance_scale,
             num_inference_steps,
+            f"sliderspace_weights/{slider_space}/{discovered_directions}",
+            slider_scale
         ],
-        outputs=[result, seed],
+        outputs=[result, slider_result, seed],
     )
 
 if __name__ == "__main__":