Update app.py

app.py

@@ -25,6 +25,24 @@ def color_loss(images, target_color=(0.1, 0.9, 0.5)):
     error = torch.abs(images - target).mean() # Mean absolute difference between the image pixels and the target color
     return error
 
+def monochromatic_loss(images, threshold=0.5, target_value=0.01):
+    # Convert images to grayscale (simple average of channels)
+    # We assume images are [N, C, H, W] where C=3 (RGB)
+    grayscale_images = (images[:,0,:,:] + images[:,1,:,:] + images[:,2,:,:]) / 3.0
+
+    # Penalize pixels that are not close to black or white
+    # Encourage values close to target_value (e.g., 0.01 for black) or 1.0 (for white)
+    # This creates a strong push towards high contrast
+    loss_black = torch.abs(grayscale_images - target_value)
+    loss_white = torch.abs(grayscale_images - (1.0 - target_value))
+
+    # For each pixel, take the minimum deviation from either black or white
+    min_deviation = torch.min(loss_black, loss_white)
+
+    # We want to minimize this deviation across the image
+    loss = min_deviation.mean()
+    return loss
+
 # And the core function to generate an image given the relevant inputs
 def generate(color, guidance_loss_scale):
   target_color = ImageColor.getcolor(color, "RGB") # Target color as RGB
@@ -36,7 +54,8 @@ def generate(color, guidance_loss_scale):
         noise_pred = image_pipe.unet(model_input, t)["sample"]
     x = x.detach().requires_grad_()
     x0 = scheduler.step(noise_pred, t, x).pred_original_sample
-    loss = color_loss(x0, target_color) * guidance_loss_scale
+    # loss = color_loss(x0, target_color) * guidance_loss_scale
+    loss  = monochromatic_loss(x0)
     cond_grad = -torch.autograd.grad(loss, x)[0]
     x = x.detach() + cond_grad 
     x = scheduler.step(noise_pred, t, x).prev_sample