bmorphism/fabric

1import functools2import random3 4import gradio as gr5import torch6 7from fabric.generator import AttentionBasedGenerator8 9 10#model_name = "dreamlike-art/dreamlike-photoreal-2.0"11model_name = ""12model_ckpt = "https://huggingface.co/Lykon/DreamShaper/blob/main/DreamShaper_7_pruned.safetensors"13 14class GeneratorWrapper:15  def __init__(self, model_name=None, model_ckpt=None):16    self.model_name = model_name if model_name else None17    self.model_ckpt = model_ckpt if model_ckpt else None18    self.dtype = torch.float16 if torch.cuda.is_available() else torch.float3219    self.device = "cuda" if torch.cuda.is_available() else "cpu"20 21    self.reload()22 23  def generate(self, *args, **kwargs):24    if not hasattr(self, "generator"):25        self.reload()26    return self.generator.generate(*args, **kwargs)27 28  def to(self, device):29    return self.generator.to(device)30 31  def reload(self):32    if hasattr(self, "generator"):33      del self.generator34    if self.device == "cuda":35      torch.cuda.empty_cache()36    self.generator = AttentionBasedGenerator(37        model_name=self.model_name,38        model_ckpt=self.model_ckpt,39        torch_dtype=self.dtype,40    ).to(self.device)41 42generator = GeneratorWrapper(model_name, model_ckpt)43 44 45css = """46.btn-green {47  background-image: linear-gradient(to bottom right, #86efac, #22c55e) !important;48  border-color: #22c55e !important;49  color: #166534 !important;50}51.btn-green:hover {52  background-image: linear-gradient(to bottom right, #86efac, #86efac) !important;53}54.btn-red {55  background: linear-gradient(to bottom right, #fda4af, #fb7185) !important;56  border-color: #fb7185 !important;57  color: #9f1239 !important;58}59.btn-red:hover {background: linear-gradient(to bottom right, #fda4af, #fda4af) !important;}60 61/*****/62 63.dark .btn-green {64  background-image: linear-gradient(to bottom right, #047857, #065f46) !important;65  border-color: #047857 !important;66  color: #ffffff !important;67}68.dark .btn-green:hover {69  background-image: linear-gradient(to bottom right, #047857, #047857) !important;70}71.dark .btn-red {72  background: linear-gradient(to bottom right, #be123c, #9f1239) !important;73  border-color: #be123c !important;74  color: #ffffff !important;75}76.dark .btn-red:hover {background: linear-gradient(to bottom right, #be123c, #be123c) !important;}77"""78 79def generate_fn(80    feedback_enabled,81    max_feedback_imgs,82    prompt,83    neg_prompt,84    liked,85    disliked,86    denoising_steps,87    guidance_scale,88    feedback_start,89    feedback_end,90    min_weight,91    max_weight,92    neg_scale,93    batch_size,94    seed,95    progress=gr.Progress(track_tqdm=True),96):97  try:98    if seed < 0:99      seed = random.randint(1,9999999999999999) #16 digits is an arbitrary limit100    print("seed: ", seed)101 102    max_feedback_imgs = max(0, int(max_feedback_imgs))103    total_images = (len(liked) if liked else 0) + (len(disliked) if disliked else 0)104 105    if not feedback_enabled:106      liked = []107      disliked = []108    elif total_images > max_feedback_imgs:109      if liked and disliked:110        max_disliked = min(len(disliked), max_feedback_imgs // 2)111        max_liked = min(len(liked), max_feedback_imgs - max_disliked)112        if max_liked > len(liked):113          max_disliked = max_feedback_imgs - max_liked114        liked = liked[-max_liked:]115        disliked = disliked[-max_disliked:]116      elif liked:117        liked = liked[-max_feedback_imgs:]118        disliked = []119      else:120        liked = []121        disliked = disliked[-max_feedback_imgs:]122    # else: keep all feedback images123    124    generate_kwargs = {125        "prompt": prompt,126        "negative_prompt": neg_prompt,127        "liked": liked,128        "disliked": disliked,129        "denoising_steps": denoising_steps,130        "guidance_scale": guidance_scale,131        "feedback_start": feedback_start,132        "feedback_end": feedback_end,133        "min_weight": min_weight,134        "max_weight": max_weight,135        "neg_scale": neg_scale,136        "seed": seed,137        "n_images": batch_size,138    }139 140    try:141      images = generator.generate(**generate_kwargs)142    except RuntimeError as err:143      if 'out of memory' in str(err):144        generator.reload()145      raise146    return [(img, f"Image {i+1}") for i, img in enumerate(images)], images, seed147  except Exception as err:148    raise gr.Error(str(err))149 150 151def add_img_from_list(i, curr_imgs, all_imgs):152  if all_imgs is None:153    all_imgs = []154  if i >= 0 and i < len(curr_imgs):155    all_imgs.append(curr_imgs[i])156  return all_imgs, all_imgs  # return (gallery, state)157 158def add_img(img, all_imgs):159  if all_imgs is None:160    all_imgs = []161  all_imgs.append(img)162  return None, all_imgs, all_imgs163 164def remove_img_from_list(event: gr.SelectData, imgs):165  if event.index >= 0 and event.index < len(imgs):166    imgs.pop(event.index)167  return imgs, imgs168 169def duplicate_seed_value(seed): #I don't like the progress bar showing on the previous seed box and this is how I hide it170    return seed171 172with gr.Blocks(css=css) as demo:173 174  liked_imgs = gr.State([])175  disliked_imgs = gr.State([])176  curr_imgs = gr.State([])177 178  with gr.Row():179    with gr.Column(scale=100):180      prompt = gr.Textbox(label="Prompt")181      neg_prompt = gr.Textbox(label="Negative prompt", value="lowres, bad anatomy, bad hands, cropped, worst quality")182    submit_btn = gr.Button("Generate", variant="primary", min_width="96px")183 184  with gr.Row(equal_height=False):185    with gr.Column():186      denoising_steps = gr.Slider(1, 100, value=20, step=1, label="Sampling steps")187      guidance_scale = gr.Slider(0.0, 30.0, value=6, step=0.25, label="CFG scale")188      batch_size = gr.Slider(1, 10, value=4, step=1, label="Batch size", interactive=False)189      seed = gr.Number(-1, minimum=-1, precision=0, label="Seed")190      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)")191      feedback_enabled = gr.Checkbox(True, label="Enable feedback", interactive=True)192 193      with gr.Accordion("Liked Images", open=True):194        liked_img_input = gr.Image(type="pil", shape=(512, 512), height=128, label="Upload liked image")195        like_gallery = gr.Gallery(label="👍 Liked images (click to remove)", columns=[3, 4, 3, 4, 5, 6], height=256, allow_preview=False)196        clear_liked_btn = gr.Button("Clear likes")197 198      with gr.Accordion("Disliked Images", open=True):199        disliked_img_input = gr.Image(type="pil", shape=(512, 512), height=128, label="Upload disliked image")200        dislike_gallery = gr.Gallery(label="👎 Disliked images (click to remove)", columns=[3, 4, 3, 4, 5, 6], height=256, allow_preview=False)201        clear_disliked_btn = gr.Button("Clear dislikes")202 203      with gr.Accordion("Feedback parameters", open=False):204        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.")205        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.")206        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)")207        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)")208        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)")209 210    with gr.Column():211      gallery = gr.Gallery(label="Generated images")212 213      like_btns = []214      dislike_btns = []215      with gr.Row():216        for i in range(0, 2):217          like_btn = gr.Button(f"👍 Image {i+1}", elem_classes="btn-green")218          like_btns.append(like_btn)219      with gr.Row():220        for i in range(2, 4):221          like_btn = gr.Button(f"👍 Image {i+1}", elem_classes="btn-green")222          like_btns.append(like_btn)223      with gr.Row():224        for i in range(0, 2):225          dislike_btn = gr.Button(f"👎 Image {i+1}", elem_classes="btn-red")226          dislike_btns.append(dislike_btn)227      with gr.Row():228        for i in range(2, 4):229          dislike_btn = gr.Button(f"👎 Image {i+1}", elem_classes="btn-red")230          dislike_btns.append(dislike_btn)231 232      prev_seed = gr.Number(-1, label="Previous seed", interactive=False)233      prev_seed_hid = gr.Number(-1, visible=False)234 235  generate_params = [236      feedback_enabled,237      max_feedback_imgs,238      prompt,239      neg_prompt,240      liked_imgs,241      disliked_imgs,242      denoising_steps,243      guidance_scale,244      feedback_start,245      feedback_end,246      feedback_min_weight,247      feedback_max_weight,248      feedback_neg_scale,249      batch_size,250      seed,251  ]252  submit_btn.click(generate_fn, generate_params, [gallery, curr_imgs, prev_seed_hid], queue=True)253  prev_seed_hid.change(duplicate_seed_value, prev_seed_hid, prev_seed, queue=False)254    255  for i, like_btn in enumerate(like_btns):256    like_btn.click(functools.partial(add_img_from_list, i), [curr_imgs, liked_imgs], [like_gallery, liked_imgs], queue=False)257  for i, dislike_btn in enumerate(dislike_btns):258    dislike_btn.click(functools.partial(add_img_from_list, i), [curr_imgs, disliked_imgs], [dislike_gallery, disliked_imgs], queue=False)259 260  like_gallery.select(remove_img_from_list, [liked_imgs], [like_gallery, liked_imgs], queue=False)261  dislike_gallery.select(remove_img_from_list, [disliked_imgs], [dislike_gallery, disliked_imgs], queue=False)262 263  liked_img_input.upload(add_img, [liked_img_input, liked_imgs], [liked_img_input, like_gallery, liked_imgs], queue=False)264  disliked_img_input.upload(add_img, [disliked_img_input, disliked_imgs], [disliked_img_input, dislike_gallery, disliked_imgs], queue=False)265 266  clear_liked_btn.click(lambda: [[], []], None, [liked_imgs, like_gallery], queue=False)267  clear_disliked_btn.click(lambda: [[], []], None, [disliked_imgs, dislike_gallery], queue=False)268 269demo.queue(1)270demo.launch(debug=True)