Update app.py

app.py

@@ -19,6 +19,8 @@ device = (
 
 pipeline_name = 'WiNE-iNEFF/Minecraft-Skin-Diffusion-V2'
 image_pipe = DDPMPipeline.from_pretrained(pipeline_name).to(device)
+scheduler = DDIMScheduler.from_pretrained(pipeline_name)
+scheduler.set_timesteps(num_inference_steps=20)
 
 def show_images_save(x):
     grid = torchvision.utils.make_grid(x, nrow=4, padding=0)
@@ -26,16 +28,14 @@ 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)):
+def color_loss(images, target_color=(0, 0, 0, 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)
-    scheduler.set_timesteps(num_inference_steps=20)
-    x = torch.randn(12, 4, 64, 64).to(device)
+    x = torch.randn(8, 4, 64, 64).to(device)
     for i, t in enumerate(scheduler.timesteps):
         model_input = scheduler.scale_model_input(x, t)
         with torch.no_grad():
@@ -104,7 +104,7 @@ with demo:
                 with gr.Row().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])
+                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(