Upload app.py

app.py

@@ -1,108 +1,127 @@
-import gradio as gr
-import torch, torchvision
-from torchvision import transforms
-import torch.nn.functional as F
-import numpy as np
-from time import time, ctime
-from PIL import Image, ImageColor
-from diffusers import DDPMPipeline
-from diffusers import DDIMScheduler
-
-device = (
-    "mps"
-    if torch.backends.mps.is_available()
-    else "cuda"
-    if torch.cuda.is_available()
-    else "cpu"
-)
-
-pipeline_name = 'WiNE-iNEFF/Minecraft-Skin-Diffusion-V2'
-image_pipe = DDPMPipeline.from_pretrained(pipeline_name, use_safetensors=True).to(device)
-try:
-    scheduler = DDIMScheduler.from_pretrained(pipeline_name)
-except:
-    scheduler = DDIMScheduler.from_pretrained(pipeline_name, subfolder="scheduler")
-scheduler.set_timesteps(num_inference_steps=25)
-
-def show_images_save(x):
-    grid = torchvision.utils.make_grid(x, nrow=4, padding=0)
-    grid_im = grid.detach().cpu().permute(1, 2, 0).clip(0, 1) * 255
-    grid_im = Image.fromarray(np.array(grid_im).astype(np.uint8))
-    #CROP
-    width, height = grid_im.size
-    sav = []
-    for w in range(width//64):
-      for h in range(height//64):
-        box = (w*64, h*64, (w+1)*64, (h+1)*64)
-        sav.append(grid_im.crop(box))
-    return sav
-
-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():
-    print(ctime(time()))
-    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():
-            noise_pred = image_pipe.unet(model_input, t)["sample"]
-        x = scheduler.step(noise_pred, t, x).prev_sample
-    return show_images_save(x)
-
-def generate_g(color, guidance_loss_scale):
-    print(f'--V2-- {ctime(time())}')
-    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 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)
-
-demo = gr.Blocks(css=".tabs {max-width: 60%; margin-left: 20%; margin-right: 20%;} #column {max-width: 50%} .fixed-height {min-height: 0px !important} @media(min-width: 1280px){.fixed-height {min-height: 0px !important}}")
-with demo:
-    gr.HTML(
-        """
-        <div style="text-align: center; margin: 0 auto;">
-            <div style="display: inline-flex; align-items: center; justify-content: center;">
-                <img src='https://huggingface.co/spaces/WiNE-iNEFF/MinecraftSkin-Diffusion/resolve/main/MSD_7.png' width='45%'>
-            </div>
-            <p style="margin-bottom: 10px; font-size: 94%; line-height: 23px;">
-                Gradio demo for Minecraft Skin Diffusion. This is simple Unconditional Diffusion Model<br>that can generate skins for game Minecraft.
-            </p>
-        </div>
-        """)
-    with gr.Tabs():
-        with gr.TabItem("V1"):
-            with gr.Column():
-                gall = gr.Gallery().style(grid=[4])
-                greet_btn = gr.Button("Generate")
-                greet_btn.click(fn=generate, outputs=gall)
-        with gr.TabItem("V2"):
-            with gr.Row():
-                with gr.Column(min_width = 0, elem_id='column'):
-                    picker = gr.ColorPicker(label="color", value="#55FFAA", min_width='0')
-                    slide = gr.Slider(label="guidance_scale", minimum=0, maximum=100, value=50)
-                    greet_btn = gr.Button("Generate", min_width = '0')
-                with gr.Column(min_width = 0):
-                  gall = gr.Gallery(min_width = '250').style(grid=[3])
-                greet_btn.click(fn=generate_g, inputs=[picker, slide], outputs=gall)
-    gr.HTML(
-        """
-        <div class="footer">
-            <div style='text-align: center;'>Minecraft Skin Diffusion by <a href='https://twitter.com/wine_ineff' target='_blank'>Artsem Holub (WiNE-iNEFF)</a></div>
-        </div>
-        """)
-
-demo.launch()
+import gradio as gr
+import torch, torchvision
+from torchvision import transforms
+import torch.nn.functional as F
+import numpy as np
+import PIL
+from PIL import Image, ImageColor
+from diffusers import DiffusionPipeline
+from diffusers import DDIMScheduler, DDPMScheduler, DEISMultistepScheduler, LCMScheduler, PNDMScheduler, UniPCMultistepScheduler
+import base64
+from urllib.parse import quote_plus
+from io import BytesIO
+
+device = ("mps" if torch.backends.mps.is_available() else "cuda" if torch.cuda.is_available() else "cpu")
+
+
+class MSPipeline(DiffusionPipeline):
+  def __init__(self, unet, scheduler):
+    super().__init__()
+    self.register_modules(unet=unet, scheduler=scheduler)
+
+  @torch.no_grad()
+  def __call__(self, batch_size = 1, num_inference_steps = 1000):
+    x = torch.randn(batch_size, self.unet.config.in_channels, self.unet.config.sample_size, self.unet.config.sample_size).to(device)
+
+    self.scheduler.set_timesteps(num_inference_steps)
+
+    for t in self.progress_bar(self.scheduler.timesteps):
+      #x = self.scheduler.scale_model_input(x, t)
+      model_output = self.unet(x, t).sample
+      x = self.scheduler.step(model_output, t, x).prev_sample
+
+    x = x.cpu().permute(0, 2, 3, 1).clip(0, 1).numpy() #* 255
+    x = self.numpy_to_pil(x)
+
+    return x
+
+
+def clear_pix(x):
+  datas = []
+  for pixel in list(x.getdata()):
+    r,g,b,a = pixel
+    if a == 0 or a < 150:
+      datas.append((0,0,0,0))
+    else:
+      datas.append((r,g,b,255))
+  x.putdata(datas)
+  return x
+
+
+def show_3D(image, print_link = False):
+  if isinstance(image, PIL.Image.Image):
+    buffer = BytesIO()
+    image.save(buffer, format="PNG")
+    skin_base64 = base64.b64encode(buffer.getvalue()).decode("utf-8")
+  else:
+    with open(image, "rb") as f:
+      skin_base64 = base64.b64encode(f.read()).decode("utf-8")
+
+  data = f"data:image/png;base64,{skin_base64}"
+  quoted = quote_plus(data)
+  url = f"https://wine-ineff.github.io/SkinViewFrame/skin.html?color=c2b4ff&data={quoted}"
+
+  if print_link == True:
+    print(url)
+  return url
+
+
+def generate(schedulers, inference_steps, images_num):
+  pipe = MSPipeline.from_pretrained("WiNE-iNEFF/Mineskin-Diffusion-v1.0", use_safetensors=True).to(device)
+
+  if schedulers == "DDIMScheduler":
+    pipe.scheduler = DDIMScheduler.from_pretrained("WiNE-iNEFF/Mineskin-Diffusion-v1.0", subfolder="scheduler")
+  elif schedulers == "DDPMScheduler":
+    pipe.scheduler = DDPMScheduler.from_pretrained("WiNE-iNEFF/Mineskin-Diffusion-v1.0", subfolder="scheduler")
+  elif schedulers == "DEISMultistepScheduler":
+    pipe.scheduler = DEISMultistepScheduler.from_pretrained("WiNE-iNEFF/Mineskin-Diffusion-v1.0", subfolder="scheduler")
+  elif schedulers == "LCMScheduler":
+    pipe.scheduler = LCMScheduler.from_pretrained("WiNE-iNEFF/Mineskin-Diffusion-v1.0", subfolder="scheduler")
+  elif schedulers == "PNDMScheduler":
+    pipe.scheduler = PNDMScheduler.from_pretrained("WiNE-iNEFF/Mineskin-Diffusion-v1.0", subfolder="scheduler")
+  elif schedulers == "UniPCMultistepScheduler":
+    pipe.scheduler = UniPCMultistepScheduler.from_pretrained("WiNE-iNEFF/Mineskin-Diffusion-v1.0", subfolder="scheduler")
+
+  images = pipe(batch_size = images_num, num_inference_steps = inference_steps)
+  return images, update_iframe(images)
+
+
+def update_iframe(images):
+  iframe_html = "<div style='display: grid; grid-template-columns: repeat(2, 1fr); gap: 10px'>"
+  for img in images:
+    iframe_url = show_3D(clear_pix(img), print_link=True)
+    iframe_html += f"<iframe style='min-width: 100%;' src='{iframe_url}'></iframe>"
+  iframe_html += "</div>"
+  return iframe_html
+
+
+interface = gr.Blocks()
+with interface:
+  gr.HTML(f"""
+    <h1 style="min-width: 100%; text-align: center;">Mineskin Diffusion</h1>
+    <p>This space run with {f"<b>{device.upper()}</b>." if device != "cpu" else f"<b>{device.upper()}</b>. This space will be work very slow. Pleasy use a <a href="https://colab.research.google.com/#fileId=https%3A//huggingface.co/WiNE-iNEFF/Mineskin-Diffusion-v1.0/MineskinDiffusion.ipynb" style="text-decoration: none; color: orange;">Google Colab notebook</a> for get fully experience of model or dublicate this space with your devices"}</p>
+  """)
+  with gr.Tabs():
+    with gr.TabItem("v1.0"):
+      with gr.Column():
+        scheduler_type = gr.Dropdown(
+            ["DDIMScheduler", "DDPMScheduler", "DEISMultistepScheduler", "LCMScheduler", "PNDMScheduler", "UniPCMultistepScheduler"],
+            label = "Type of scheduler"
+        )
+
+        with gr.Row():
+          inference_steps = gr.Number(
+              label = "Amount of denoising steps",
+              value = 30, minimum = 5, maximum = 1000
+          )
+          images_num = gr.Number(
+              label = "Amount of skins",
+              value = 4, minimum = 1, maximum = 4
+          )
+
+        gallery = gr.Gallery(columns=4, object_fit="scale-down", min_width=76)
+        iframe_output = gr.HTML()
+        gen_btn = gr.Button("Generate")
+        gen_btn.click(fn=generate, outputs=[gallery, iframe_output], inputs=[scheduler_type, inference_steps, images_num])
+
+interface.launch()