Updates for spaces

app.py

@@ -15,11 +15,9 @@ with gr.Blocks() as demo:
         ReSWD is a method for distribution matching with reduced variance.
         """
     )
-    fabric = L.Fabric(devices=1, accelerator="auto", precision="16-mixed")
-
     with gr.Tab("SW Guidance (SD 3.5 Large Turbo)"):
-        create_sw_guidance(fabric, "stabilityai/stable-diffusion-3.5-large-turbo")
+        create_sw_guidance("stabilityai/stable-diffusion-3.5-large-turbo")
     with gr.Tab("Color Matching"):
-        create_color_matching(fabric)
+        create_color_matching()
 
     demo.launch()

src/color_matcher.py

@@ -61,7 +61,6 @@ class CDL(torch.nn.Module):
 
 
 def train(
-    fabric: L.Fabric,
     criteria: AbstractLoss,
     source_img: Float[torch.Tensor, "B C H W"],
     target_img: Float[torch.Tensor, "B C H W"],
@@ -71,25 +70,25 @@ def train(
     silent: bool = False,
     write_video_animation_path: Optional[str] = None,
 ) -> Tuple[Float[torch.Tensor, "*B C H W"], CDL, List[float]]:
-    criteria = fabric.setup(criteria)
+    criteria = criteria.cuda()
 
     source_max_res = Resize(match_resolution, antialias=True)(source_img)
     target_max_res = Resize(match_resolution, antialias=True)(target_img)
 
     target_cielab = (
-        fabric.to_device(rgb_to_lab(target_max_res).permute(0, 3, 1, 2))
+        rgb_to_lab(target_max_res).cuda().permute(0, 3, 1, 2)
         .permute(0, 2, 3, 1)
         .contiguous()
     )
 
-    source_max_res = fabric.to_device(source_max_res)
-    source_img = fabric.to_device(source_img)
+    source_max_res = source_max_res.cuda()
+    source_img = source_img.cuda()
 
     batch_size = source_img.shape[0]
     cdl = CDL(batch_size)
 
     optim = torch.optim.Adam(cdl.parameters(), lr=lr)
-    cdl, optim = fabric.setup(cdl, optim)
+    cdl, optim = cdl.cuda(), optim.cuda()
 
     lossses = []
     for i in tqdm(range(num_steps), disable=silent):
@@ -106,7 +105,7 @@ def train(
             i,
         )
 
-        fabric.backward(loss)
+        loss.backward()
         optim.step()
 
         lossses.append(loss.item())

src/gradio_demo/color_matching.py

@@ -10,7 +10,7 @@ from src.loss import VectorSWDLoss
 from src.utils.image import from_torch, to_torch
 
 
-def create_color_matching(fabric: L.Fabric):
+def create_color_matching():
     """
     Creates the Gradio interface for color matching between source and target images.
     """

src/gradio_demo/sw_guidance.py

@@ -44,7 +44,7 @@ def log_slider_to_lr(log_lr):
 
 
 def create_sw_guidance(
-    fabric: L.Fabric, model_name: str = "stabilityai/stable-diffusion-3.5-large"
+    model_name: str = "stabilityai/stable-diffusion-3.5-large"
 ):
     """
     Creates the Gradio interface for SW guidance with SD3.5.
@@ -63,7 +63,7 @@ def create_sw_guidance(
     pipe = create_pipeline(
         model_name,
         device="cuda",
-        compile=True,
+        compile=False,
     )
 
     model_config = models[model_name]