add sd

app.py

@@ -414,8 +414,11 @@ def run_fn(
     images = [tup[0] for tup in images]
     images = [transform_image(image, resolution=resolution) for image in images]
     images = torch.stack(images)
-        
+    
     model = load_model(model_name)
+    if "stable" in model_name.lower() and "diffusion" in model_name.lower():
+        model.timestep = layer
+        layer = 1
     
     kwargs = {
         "model_name": model_name,
@@ -593,7 +596,8 @@ def make_output_images_section():
 
 def make_parameters_section():
     gr.Markdown("### Parameters <a style='color: #0044CC;' href='https://ncut-pytorch.readthedocs.io/en/latest/how_to_get_better_segmentation/' target='_blank'>Help</a>")
-    from ncut_pytorch.backbone import list_models, get_demo_model_names
+    # from ncut_pytorch.backbone import list_models, get_demo_model_names
+    from backbone import list_models, get_demo_model_names
     model_names = list_models()
     model_names = sorted(model_names)
     model_dropdown = gr.Dropdown(model_names, label="Backbone", value="DiNO(dino_vitb8_448)", elem_id="model_name")
@@ -602,6 +606,13 @@ def make_parameters_section():
     num_eig_slider = gr.Slider(1, 1000, step=1, label="NCUT: Number of eigenvectors", value=100, elem_id="num_eig", info='increase for more clusters')
 
     def change_layer_slider(model_name):
+        # SD2, UNET
+        if "stable" in model_name.lower() and "diffusion" in model_name.lower():
+            from backbone import SD_KEY_DICT
+            default_layer = 'up_2_resnets_1_block' if 'diffusion-3' not in model_name else 'block_23'
+            return (gr.Slider(1, 49, step=1, label="Diffusion: Timestep (Noise)", value=5, elem_id="layer", visible=True, info="Noise level, 50 is max noise"),
+                    gr.Dropdown(SD_KEY_DICT[model_name], label="Diffusion: Layer and Node", value=default_layer, elem_id="node_type", info="From the SD U-Net"))
+        
         layer_dict = LAYER_DICT
         if model_name in layer_dict:
             value = layer_dict[model_name]
@@ -657,6 +668,8 @@ with demo:
                 logging_text = gr.Textbox("Logging information", label="Logging", elem_id="logging", type="text", placeholder="Logging information")
 
         clear_images_button.click(lambda x: ([], []), outputs=[input_gallery, output_gallery])
+        
+        
         submit_button.click(
             run_fn,
             inputs=[
@@ -917,3 +930,6 @@ if DOWNLOAD_ALL_MODELS_DATASETS:
     threading.Thread(target=download_all_datasets).start()
     
 demo.launch(share=True)
+
+
+# %%

docker-compose.yml

@@ -0,0 +1,22 @@
+version: '3.8'
+
+services:
+  ncut_demo:
+    image: registry.hf.space/huzey-ncut-pytorch:latest
+    container_name: ncut_demo
+    environment:
+      - HF_ACCESS_TOKEN=${HF_ACCESS_TOKEN}
+      - USE_HUGGINGFACE_ZEROGPU=false
+    ports:
+      - "7860:7860"
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - capabilities: [gpu]
+    platform: linux/amd64
+    runtime: nvidia
+    shm_size: '64G'
+    command: python app.py
+    networks:
+      - default

requirements.txt

@@ -1,6 +1,6 @@
 torch
 torchvision 
-ncut-pytorch>=1.2.7
+ncut-pytorch>=1.3.1
 opencv-python
 decord
 transformers