import functools 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 = None 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 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 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) 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], queue=True) 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)