Update app.py

app.py

@@ -89,31 +89,31 @@ MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
 
-base_model_id = "black-forest-labs/FLUX.1-schnell"
-max_sequence_length = 256
-flux_pipe = FluxPipeline.from_pretrained(base_model_id, torch_dtype=torch_dtype)
-flux_pipe.vae = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=torch_dtype)
-flux_pipe = flux_pipe.to(device).to(torch_dtype)
-# pipe.enable_sequential_cpu_offload()
-transformer = flux_pipe.transformer
+# base_model_id = "black-forest-labs/FLUX.1-schnell"
+# max_sequence_length = 256
+# flux_pipe = FluxPipeline.from_pretrained(base_model_id, torch_dtype=torch_dtype)
+# flux_pipe.vae = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=torch_dtype)
+# flux_pipe = flux_pipe.to(device).to(torch_dtype)
+# # pipe.enable_sequential_cpu_offload()
+# transformer = flux_pipe.transformer
 
-## Change these parameters based on how you trained your sliderspace sliders
-train_method = 'flux-attn'
-rank = 1 
-alpha =1 
+# ## Change these parameters based on how you trained your sliderspace sliders
+# train_method = 'flux-attn'
+# rank = 1 
+# alpha =1 
 
-flux_networks = {}
-modules = DEFAULT_TARGET_REPLACE
-modules += UNET_TARGET_REPLACE_MODULE_CONV
-for i in range(1):
-    flux_networks[i] = LoRANetwork(
-        transformer,
-        rank=int(rank),
-        multiplier=1.0,
-        alpha=int(alpha),
-        train_method=train_method,
-        fast_init=True,
-    ).to(device, dtype=torch_dtype)
+# flux_networks = {}
+# modules = DEFAULT_TARGET_REPLACE
+# modules += UNET_TARGET_REPLACE_MODULE_CONV
+# for i in range(1):
+#     flux_networks[i] = LoRANetwork(
+#         transformer,
+#         rank=int(rank),
+#         multiplier=1.0,
+#         alpha=int(alpha),
+#         train_method=train_method,
+#         fast_init=True,
+#     ).to(device, dtype=torch_dtype)
 
 
 def update_sliderspace_choices(model_choice):
@@ -175,38 +175,38 @@ def infer(
                 height=height,
                 generator=generator,
             ).images[0]
-    else:
-        sliderspace_path = f"flux_sliderspace_weights/{slider_space}/slider_{int(discovered_directions.split(' ')[-1])-1}.pt"
-        for net in flux_networks:
-            flux_networks[net].load_state_dict(torch.load(sliderspace_path))
-            flux_networks[net].set_lora_slider(-1*slider_scale)
-        with flux_networks[0]:
-            pass
-
-        # original image
-        generator = torch.Generator().manual_seed(seed)
-        image = flux_pipe(
-                prompt=prompt,
-                guidance_scale=guidance_scale,
-                num_inference_steps=num_inference_steps,
-                width=width,
-                height=height,
-                generator=generator,
-                max_sequence_length = 256,
-            ).images[0]
+    # else:
+    #     sliderspace_path = f"flux_sliderspace_weights/{slider_space}/slider_{int(discovered_directions.split(' ')[-1])-1}.pt"
+    #     for net in flux_networks:
+    #         flux_networks[net].load_state_dict(torch.load(sliderspace_path))
+    #         flux_networks[net].set_lora_slider(-1*slider_scale)
+    #     with flux_networks[0]:
+    #         pass
+
+    #     # original image
+    #     generator = torch.Generator().manual_seed(seed)
+    #     image = flux_pipe(
+    #             prompt=prompt,
+    #             guidance_scale=guidance_scale,
+    #             num_inference_steps=num_inference_steps,
+    #             width=width,
+    #             height=height,
+    #             generator=generator,
+    #             max_sequence_length = 256,
+    #         ).images[0]
     
-        # edited image
-        generator = torch.Generator().manual_seed(seed)
-        with flux_networks[0]:
-            slider_image = flux_pipe(
-                prompt=prompt,
-                guidance_scale=guidance_scale,
-                num_inference_steps=num_inference_steps,
-                width=width,
-                height=height,
-                generator=generator,
-                max_sequence_length = 256,
-            ).images[0]
+    #     # edited image
+    #     generator = torch.Generator().manual_seed(seed)
+    #     with flux_networks[0]:
+    #         slider_image = flux_pipe(
+    #             prompt=prompt,
+    #             guidance_scale=guidance_scale,
+    #             num_inference_steps=num_inference_steps,
+    #             width=width,
+    #             height=height,
+    #             generator=generator,
+    #             max_sequence_length = 256,
+    #         ).images[0]
     
     return image, slider_image, seed
 
@@ -280,7 +280,7 @@ with gr.Blocks(css=css) as demo:
 
         # Add model selection dropdown
         model_choice = gr.Dropdown(
-            choices=["SDXL-DMD", "FLUX-Schnell"],
+            choices=["SDXL-DMD"],
             label="Model",
             value="SDXL-DMD"
         )