Update app.py

app.py

@@ -26,6 +26,11 @@ def show_images_save(x):
     grid_im = Image.fromarray(np.array(grid_im).astype(np.uint8))
     return grid_im
 
+def color_loss(images, target_color=(0.1, 0.9, 0.5, 1)):
+    target = (torch.tensor(target_color).to(images.device) * 2 - 1)
+    target = target[None, :, None, None]
+    error = torch.abs(images - target).mean()
+    return error
 
 def generate():
     scheduler = DDIMScheduler.from_pretrained(pipeline_name)
@@ -38,6 +43,21 @@ def generate():
         x = scheduler.step(noise_pred, t, x).prev_sample
     return show_images_save(x)
 
+def generate_g(color, guidance_loss_scale):
+    target_color = ImageColor.getcolor(color, "RGBA")  # Target color as RGB
+    target_color = [a / 255 for a in target_color]  # Rescale from (0, 255) to (0, 1)
+    x = torch.randn(8, 4, 64, 64).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
+    return show_images_save(x)
 
 def crrop(file):
     width, height = file.size
@@ -48,14 +68,17 @@ def crrop(file):
         sav.append(file.crop(box))
     return sav
 
-
 def ex():
     t = time()
     print(ctime(t))
     return crrop(generate())
 
+def ex_g(picker, slider):
+    t = time()
+    print(ctime(t))
+    return crrop(generate_g(picker, slider))
 
-demo = gr.Blocks(css=".container {max-width: 730px;margin: auto;} #gallery {margin-bottom: 15px; margin-left: auto; margin-right: auto; border-bottom-right-radius: .5rem !important; border-bottom-left-radius: .5rem !important;}")
+demo = gr.Blocks(css="#row-1 {max-height: 300px !important;} .container {max-width: 730px;margin: auto;} #gallery {margin-bottom: 15px; margin-left: auto; margin-right: auto; border-bottom-right-radius: .5rem !important; border-bottom-left-radius: .5rem !important;}")
 
 with demo:
     gr.HTML(
@@ -72,10 +95,18 @@ with demo:
     with gr.Tabs():
         with gr.TabItem("V1"):
             with gr.Column():
-                with gr.Row().style(equal_height=True):
+                with gr.Row(elem_id='row-1').style(equal_height=True):
                     gall = gr.Gallery(elem_id='gallery').style(grid=[4])
                 greet_btn = gr.Button("Generate")
                 greet_btn.click(fn=ex, outputs=gall)
+        with gr.TabItem("V2"):
+            with gr.Column():
+                with gr.Row(elem_id='row-1').style(equal_height=True):
+                    picker = gr.ColorPicker(label="color", value="#55FFAA")
+                    slide = gr.Slider(label="guidance_scale", minimum=0, maximum=100, value=50)
+                    gall = gr.Gallery(elem_id='gallery').style(grid=[4])
+                greet_btn = gr.Button("Generate")
+                greet_btn.click(fn=ex_g, inputs=[picker, slide] outputs=gall)
     gr.HTML(
         """
         <div class="footer">