Update app.py

app.py

@@ -1,5 +1,6 @@
 import gradio as gr
-import torch, torchvision
+import torch
+import torchvision
 import torch.nn.functional as F
 import numpy as np
 from PIL import Image, ImageColor
@@ -27,24 +28,25 @@ def color_loss(images, target_color=(0.1, 0.9, 0.5)):
 
 # 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
-  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 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
+    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 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 = (im * 255).byte().numpy()  # Convert to uint8 numpy array
+    im = Image.fromarray(im)
+    im.save('test.jpeg')
+    return im
 
 # See the gradio docs for the types of inputs and outputs available
 inputs = [
@@ -65,5 +67,4 @@ demo = gr.Interface(
 
 # And launching
 if __name__ == "__main__":
-    demo.launch(enable_queue=True)
-
+    demo.launch()  # Removed enable_queue=True