Update app.py

app.py

@@ -14,7 +14,7 @@ image_pipe = DDPMPipeline.from_pretrained(pipeline_name).to(device)
 
 # Set up the scheduler
 scheduler = DDIMScheduler.from_pretrained(pipeline_name)
-scheduler.set_timesteps(num_inference_steps=40)
+scheduler.set_timesteps(num_inference_steps=20)
 
 def color_loss(images, target_color=(0.1, 0.9, 0.5)):
     """Given a target color (R, G, B) return a loss for how far away on average 
@@ -26,49 +26,24 @@ def color_loss(images, target_color=(0.1, 0.9, 0.5)):
 
 
 def generate(color, guidance_loss_scale):
-    
-    # Target color as RGB
-    target_color = ImageColor.getcolor(color, "RGB")
-
-    # Rescale from (0, 255) to (0, 1)
-    target_color = [a/255 for a in target_color] 
-
-    # Initial random x - just one image but you could add a 'num_images' argument/input to give the user control
-    x = torch.randn(1, 3, 256, 256).to(device) 
-
-    # Our custom sampling loop:
-    for i, t in tqdm(enumerate(scheduler.timesteps)):
-        
-        # Prep the model input 
-        model_input = scheduler.scale_model_input(x, t)
-    
-        # predict the noise residual
-        with torch.no_grad():
-            noise_pred = image_pipe.unet(model_input, t)["sample"]
-    
-        # Set requires grad on x (shortcut method - we're doing this AFTER the unet)
-        x = x.detach().requires_grad_()
-    
-        # Get the predicted x0:
-        x0 = scheduler.step(noise_pred, t, x).pred_original_sample
-    
-        # Calculate loss
-        loss = color_loss(x0, target_color) * guidance_loss_scale
-    
-        # Get gradient
-        cond_grad = -torch.autograd.grad(loss, x)[0]
-    
-        # Modify x based on this gradient
-        x = x.detach() + cond_grad 
-        
-        # Now step with scheduler
-        x = scheduler.step(noise_pred, t, x).prev_sample
-        
-    # Return the final output as an image (or image grid if there are more than one images)
-    grid = torchvision.utils.make_grid(x, nrow=4)
-    im = grid.permute(1, 2, 0).cpu().clip(-1, 1)*0.5 + 0.5
-    return Image.fromarray(np.array(im*255).astype(np.uint8))
-
+  target_color = ImageColor.getcolor(color, "RGB") # Target color as RGB
+  target_color = [a/255 for a in target_color] # Rescale from (0, 255) to (0, 1)
+  x = torch.randn(1, 3, 256, 256).to(device) 
+  for i, t in tqdm(enumerate(scheduler.timesteps)):
+    model_input = scheduler.scale_model_input(x, t)
+    with torch.no_grad():
+        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
+    cond_grad = -torch.autograd.grad(loss, x)[0]
+    x = x.detach() + cond_grad 
+    x = scheduler.step(noise_pred, t, x).prev_sample
+  grid = torchvision.utils.make_grid(x, nrow=4)
+  im = grid.permute(1, 2, 0).cpu().clip(-1, 1)*0.5 + 0.5
+  im = Image.fromarray(np.array(im*255).astype(np.uint8))
+  im.save('test.jpeg')
+  return im
 
 inputs = [
     gr.ColorPicker(label="color", value='55FFAA'), # Add any inputs you need here
@@ -83,6 +58,7 @@ demo = gr.Interface(
     examples=[
         ["#BB2266"],["#44CCAA"] # You can provide some example inputs to get people started
     ],
+)
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(enable_queue=True)