Init

app.py

@@ -0,0 +1,268 @@
+import functools
+import random
+
+import gradio as gr
+import torch
+
+from fabric.generator import AttentionBasedGenerator
+
+
+#model_name = "dreamlike-art/dreamlike-photoreal-2.0"
+model_name = ""
+model_ckpt = "https://huggingface.co/Lykon/DreamShaper/blob/main/DreamShaper_7_pruned.safetensors"
+
+class GeneratorWrapper:
+  def __init__(self, model_name=None, model_ckpt=None):
+    self.model_name = model_name if model_name else None
+    self.model_ckpt = model_ckpt if model_ckpt else None
+    self.dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+    self.device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    self.reload()
+
+  def generate(self, *args, **kwargs):
+    if not hasattr(self, "generator"):
+        self.reload()
+    return self.generator.generate(*args, **kwargs)
+
+  def to(self, device):
+    return self.generator.to(device)
+
+  def reload(self):
+    if hasattr(self, "generator"):
+      del self.generator
+    if self.device == "cuda":
+      torch.cuda.empty_cache()
+    self.generator = AttentionBasedGenerator(
+        model_name=self.model_name,
+        model_ckpt=self.model_ckpt,
+        torch_dtype=self.dtype,
+    ).to(self.device)
+
+generator = GeneratorWrapper(model_name, model_ckpt)
+
+
+css = """
+.btn-green {
+  background-image: linear-gradient(to bottom right, #86efac, #22c55e) !important;
+  border-color: #22c55e !important;
+  color: #166534 !important;
+}
+.btn-green:hover {
+  background-image: linear-gradient(to bottom right, #86efac, #86efac) !important;
+}
+.btn-red {
+  background: linear-gradient(to bottom right, #fda4af, #fb7185) !important;
+  border-color: #fb7185 !important;
+  color: #9f1239 !important;
+}
+.btn-red:hover {background: linear-gradient(to bottom right, #fda4af, #fda4af) !important;}
+/*****/
+.dark .btn-green {
+  background-image: linear-gradient(to bottom right, #047857, #065f46) !important;
+  border-color: #047857 !important;
+  color: #ffffff !important;
+}
+.dark .btn-green:hover {
+  background-image: linear-gradient(to bottom right, #047857, #047857) !important;
+}
+.dark .btn-red {
+  background: linear-gradient(to bottom right, #be123c, #9f1239) !important;
+  border-color: #be123c !important;
+  color: #ffffff !important;
+}
+.dark .btn-red:hover {background: linear-gradient(to bottom right, #be123c, #be123c) !important;}
+"""
+
+def generate_fn(
+    feedback_enabled,
+    max_feedback_imgs,
+    prompt,
+    neg_prompt,
+    liked,
+    disliked,
+    denoising_steps,
+    guidance_scale,
+    feedback_start,
+    feedback_end,
+    min_weight,
+    max_weight,
+    neg_scale,
+    batch_size,
+    seed,
+    progress=gr.Progress(track_tqdm=True),
+):
+  try:
+    if seed < 0:
+      seed = random.randint(1,9999999999999999) #16 digits is an arbitrary limit
+    print("seed: ", seed)
+
+    max_feedback_imgs = max(0, int(max_feedback_imgs))
+    total_images = (len(liked) if liked else 0) + (len(disliked) if disliked else 0)
+
+    if not feedback_enabled:
+      liked = []
+      disliked = []
+    elif total_images > max_feedback_imgs:
+      if liked and disliked:
+        max_disliked = min(len(disliked), max_feedback_imgs // 2)
+        max_liked = min(len(liked), max_feedback_imgs - max_disliked)
+        if max_liked > len(liked):
+          max_disliked = max_feedback_imgs - max_liked
+        liked = liked[-max_liked:]
+        disliked = disliked[-max_disliked:]
+      elif liked:
+        liked = liked[-max_feedback_imgs:]
+        disliked = []
+      else:
+        liked = []
+        disliked = disliked[-max_feedback_imgs:]
+    # else: keep all feedback images
+    
+    generate_kwargs = {
+        "prompt": prompt,
+        "negative_prompt": neg_prompt,
+        "liked": liked,
+        "disliked": disliked,
+        "denoising_steps": denoising_steps,
+        "guidance_scale": guidance_scale,
+        "feedback_start": feedback_start,
+        "feedback_end": feedback_end,
+        "min_weight": min_weight,
+        "max_weight": max_weight,
+        "neg_scale": neg_scale,
+        "seed": seed,
+        "n_images": batch_size,
+    }
+
+    try:
+      images = generator.generate(**generate_kwargs)
+    except RuntimeError as err:
+      if 'out of memory' in str(err):
+        generator.reload()
+      raise
+    return [(img, f"Image {i+1}") for i, img in enumerate(images)], images, seed
+  except Exception as err:
+    raise gr.Error(str(err))
+
+
+def add_img_from_list(i, curr_imgs, all_imgs):
+  if all_imgs is None:
+    all_imgs = []
+  if i >= 0 and i < len(curr_imgs):
+    all_imgs.append(curr_imgs[i])
+  return all_imgs, all_imgs  # return (gallery, state)
+
+def add_img(img, all_imgs):
+  if all_imgs is None:
+    all_imgs = []
+  all_imgs.append(img)
+  return None, all_imgs, all_imgs
+
+def remove_img_from_list(event: gr.SelectData, imgs):
+  if event.index >= 0 and event.index < len(imgs):
+    imgs.pop(event.index)
+  return imgs, imgs
+
+def duplicate_seed_value(seed): #I don't like the progress bar showing on the previous seed box and this is how I hide it
+    return seed
+
+with gr.Blocks(css=css) as demo:
+
+  liked_imgs = gr.State([])
+  disliked_imgs = gr.State([])
+  curr_imgs = gr.State([])
+
+  with gr.Row():
+    with gr.Column(scale=100):
+      prompt = gr.Textbox(label="Prompt")
+      neg_prompt = gr.Textbox(label="Negative prompt", value="lowres, bad anatomy, bad hands, cropped, worst quality")
+    submit_btn = gr.Button("Generate", variant="primary", min_width="96px")
+
+  with gr.Row(equal_height=False):
+    with gr.Column():
+      denoising_steps = gr.Slider(1, 100, value=20, step=1, label="Sampling steps")
+      guidance_scale = gr.Slider(0.0, 30.0, value=6, step=0.25, label="CFG scale")
+      batch_size = gr.Slider(1, 10, value=4, step=1, label="Batch size", interactive=False)
+      seed = gr.Number(-1, minimum=-1, precision=0, label="Seed")
+      max_feedback_imgs = gr.Slider(0, 20, value=6, step=1, label="Max. feedback images", info="Maximum number of liked/disliked images to be used. If exceeded, only the most recent images will be used as feedback. (NOTE: large number of feedback imgs => high VRAM requirements)")
+      feedback_enabled = gr.Checkbox(True, label="Enable feedback", interactive=True)
+
+      with gr.Accordion("Liked Images", open=True):
+        liked_img_input = gr.Image(type="pil", shape=(512, 512), height=128, label="Upload liked image")
+        like_gallery = gr.Gallery(label="👍 Liked images (click to remove)", columns=[3, 4, 3, 4, 5, 6], height=256, allow_preview=False)
+        clear_liked_btn = gr.Button("Clear likes")
+
+      with gr.Accordion("Disliked Images", open=True):
+        disliked_img_input = gr.Image(type="pil", shape=(512, 512), height=128, label="Upload disliked image")
+        dislike_gallery = gr.Gallery(label="👎 Disliked images (click to remove)", columns=[3, 4, 3, 4, 5, 6], height=256, allow_preview=False)
+        clear_disliked_btn = gr.Button("Clear dislikes")
+
+      with gr.Accordion("Feedback parameters", open=False):
+        feedback_start = gr.Slider(0.0, 1.0, value=0.0, label="Feedback start", info="Fraction of denoising steps starting from which to use max. feedback weight.")
+        feedback_end = gr.Slider(0.0, 1.0, value=0.8, label="Feedback end", info="Up to what fraction of denoising steps to use max. feedback weight.")
+        feedback_min_weight = gr.Slider(0.0, 1.0, value=0.0, label="Feedback min. weight", info="Attention weight of feedback images when turned off (set to 0.0 to disable)")
+        feedback_max_weight = gr.Slider(0.0, 1.0, value=0.8, label="Feedback max. weight", info="Attention weight of feedback images when turned on (set to 0.0 to disable)")
+        feedback_neg_scale = gr.Slider(0.0, 1.0, value=0.5, label="Neg. feedback scale", info="Attention weight of disliked images relative to liked images (set to 0.0 to disable negative feedback)")
+
+    with gr.Column():
+      gallery = gr.Gallery(label="Generated images")
+
+      like_btns = []
+      dislike_btns = []
+      with gr.Row():
+        for i in range(0, 2):
+          like_btn = gr.Button(f"👍 Image {i+1}", elem_classes="btn-green")
+          like_btns.append(like_btn)
+      with gr.Row():
+        for i in range(2, 4):
+          like_btn = gr.Button(f"👍 Image {i+1}", elem_classes="btn-green")
+          like_btns.append(like_btn)
+      with gr.Row():
+        for i in range(0, 2):
+          dislike_btn = gr.Button(f"👎 Image {i+1}", elem_classes="btn-red")
+          dislike_btns.append(dislike_btn)
+      with gr.Row():
+        for i in range(2, 4):
+          dislike_btn = gr.Button(f"👎 Image {i+1}", elem_classes="btn-red")
+          dislike_btns.append(dislike_btn)
+
+      prev_seed = gr.Number(-1, label="Previous seed", interactive=False)
+      prev_seed_hid = gr.Number(-1, visible=False)
+
+  generate_params = [
+      feedback_enabled,
+      max_feedback_imgs,
+      prompt,
+      neg_prompt,
+      liked_imgs,
+      disliked_imgs,
+      denoising_steps,
+      guidance_scale,
+      feedback_start,
+      feedback_end,
+      feedback_min_weight,
+      feedback_max_weight,
+      feedback_neg_scale,
+      batch_size,
+      seed,
+  ]
+  submit_btn.click(generate_fn, generate_params, [gallery, curr_imgs, prev_seed_hid], queue=True)
+  prev_seed_hid.change(duplicate_seed_value, prev_seed_hid, prev_seed, queue=False)
+    
+  for i, like_btn in enumerate(like_btns):
+    like_btn.click(functools.partial(add_img_from_list, i), [curr_imgs, liked_imgs], [like_gallery, liked_imgs], queue=False)
+  for i, dislike_btn in enumerate(dislike_btns):
+    dislike_btn.click(functools.partial(add_img_from_list, i), [curr_imgs, disliked_imgs], [dislike_gallery, disliked_imgs], queue=False)
+
+  like_gallery.select(remove_img_from_list, [liked_imgs], [like_gallery, liked_imgs], queue=False)
+  dislike_gallery.select(remove_img_from_list, [disliked_imgs], [dislike_gallery, disliked_imgs], queue=False)
+
+  liked_img_input.upload(add_img, [liked_img_input, liked_imgs], [liked_img_input, like_gallery, liked_imgs], queue=False)
+  disliked_img_input.upload(add_img, [disliked_img_input, disliked_imgs], [disliked_img_input, dislike_gallery, disliked_imgs], queue=False)
+
+  clear_liked_btn.click(lambda: [[], []], None, [liked_imgs, like_gallery], queue=False)
+  clear_disliked_btn.click(lambda: [[], []], None, [disliked_imgs, dislike_gallery], queue=False)
+
+demo.queue(1)
+demo.launch(debug=True)