Create app.py

app.py

@@ -0,0 +1,84 @@
+import gradio as gr
+import torch, torchvision
+import torch.nn.functional as F
+import numpy as np
+from PIL import Image, ImageColor
+from diffusers import DDPMPipeline
+from diffusers import DDIMScheduler
+
+# Load the pretrained pipeline
+pipeline_name = 'johnowhitaker/sd-class-wikiart-from-bedrooms'
+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)
+
+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 
+    the images' pixels are from that color. Defaults to a light teal: (0.1, 0.9, 0.5) """
+    target = torch.tensor(target_color).to(images.device) * 2 - 1 # Map target color to (-1, 1)
+    target = target[None, :, None, None] # Get shape right to work with the images (b, c, h, w)
+    error = torch.abs(images - target).mean() # Mean absolute difference between the image pixels and the target color
+    return error
+
+
+def generate(color, guidance_loss_scale):
+    
+    # Target color as RGB
+    target_color = ImageColor.getcolor(color, "RGB")
+
+    # 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))
+
+
+inputs = [
+    gr.ColorPicker(label="color", value='55FFAA'), # Add any inputs you need here
+    gr.Slider(label="guidance_scale", minimum=1, maximum=100, value=30)
+]
+outputs = gr.Image(label="result")
+
+demo = gr.Interface(
+    fn=generate,
+    inputs=inputs,
+    outputs=outputs,
+    examples=[
+        ["#BB2266"], # You can provide some example inputs to get people started
+    ],
+)
+
+if __name__ == "__main__":
+    demo.launch()