Update app.py

app.py

@@ -1,75 +1,129 @@
-import gradio as gr
 import os
-import shutil
-from main import fine_tune_model
-from diffusers import StableDiffusionPipeline, DDIMScheduler
-import torch
 
-MODEL_NAME = "runwayml/stable-diffusion-v1-5"
-OUTPUT_DIR = "/home/user/app/stable_diffusion_weights/custom_model"
+import huggingface_hub, spaces
+huggingface_hub.snapshot_download(repo_id='tsujuifu/ml-mgie', repo_type='model', local_dir='_ckpt', local_dir_use_symlinks=False)
+os.system('ls _ckpt')
 
-def fine_tune(instance_prompt, image1, image2=None):
-    instance_data_dir = "/home/user/app/instance_images"
-    
-    try:
-        if os.path.exists(instance_data_dir):
-            shutil.rmtree(instance_data_dir)
-        os.makedirs(instance_data_dir, exist_ok=True)
-        
-        image1.save(os.path.join(instance_data_dir, "instance_0.png"))
-        if image2 is not None:
-            image2.save(os.path.join(instance_data_dir, "instance_1.png"))
-        
-        fine_tune_model(instance_data_dir, instance_prompt, MODEL_NAME, OUTPUT_DIR)
-        return "Model fine-tuning complete."
-    except Exception as e:
-        return str(e)
-
-def generate_images(prompt, num_samples, height, width, num_inference_steps, guidance_scale):
-    try:
-        if not os.path.exists(OUTPUT_DIR):
-            return "The model path does not exist."
-        
-        pipe = StableDiffusionPipeline.from_pretrained(OUTPUT_DIR, safety_checker=None, torch_dtype=torch.float32).to("cpu")
-        pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
-        
-        with torch.autocast("cpu"), torch.inference_mode():
-            images = pipe(
-                prompt, height=height, width=width, num_images_per_prompt=num_samples,
-                num_inference_steps=num_inference_steps, guidance_scale=guidance_scale
-            ).images
-        
-        return images
-    except Exception as e:
-        return str(e)
-
-def gradio_app():
-    with gr.Blocks() as demo:
-        with gr.Tab("Fine-Tune Model"):
-            with gr.Row():
-                with gr.Column():
-                    instance_prompt = gr.Textbox(label="Instance Prompt")
-                    image1 = gr.Image(label="Upload Image 1", type="pil")
-                    image2 = gr.Image(label="Upload Image 2 (Optional)", type="pil")
-                    fine_tune_button = gr.Button("Fine-Tune Model")
-                    output_text = gr.Textbox(label="Output")
-                fine_tune_button.click(fine_tune, inputs=[instance_prompt, image1, image2], outputs=output_text)
+from PIL import Image
+
+import numpy as np
+import torch as T
+import transformers, diffusers
+
+from conversation import conv_templates
+from mgie_llava import *
+
+import gradio as gr
+
+def crop_resize(f, sz=512):
+    w, h = f.size
+    if w>h:
+        p = (w-h)//2
+        f = f.crop([p, 0, p+h, h])
+    elif h>w:
+        p = (h-w)//2
+        f = f.crop([0, p, w, p+w])
+    f = f.resize([sz, sz])
+    return f
+def remove_alter(s):  # hack expressive instruction
+    if 'ASSISTANT:' in s: s = s[s.index('ASSISTANT:')+10:].strip()
+    if '</s>' in s: s = s[:s.index('</s>')].strip()
+    if 'alternative' in s.lower(): s = s[:s.lower().index('alternative')]
+    if '[IMG0]' in s: s = s[:s.index('[IMG0]')]
+    s = '.'.join([s.strip() for s in s.split('.')[:2]])
+    if s[-1]!='.': s += '.'
+    return s.strip()
+
+DEFAULT_IMAGE_TOKEN = '<image>'
+DEFAULT_IMAGE_PATCH_TOKEN = '<im_patch>'
+DEFAULT_IM_START_TOKEN = '<im_start>'
+DEFAULT_IM_END_TOKEN = '<im_end>'
+PATH_LLAVA = '_ckpt/LLaVA-7B-v1'
+
+tokenizer = transformers.AutoTokenizer.from_pretrained(PATH_LLAVA)
+model = LlavaLlamaForCausalLM.from_pretrained(PATH_LLAVA, low_cpu_mem_usage=True, torch_dtype=T.float16, use_cache=True).cuda()
+image_processor = transformers.CLIPImageProcessor.from_pretrained(model.config.mm_vision_tower, torch_dtype=T.float16)
+
+tokenizer.padding_side = 'left'
+tokenizer.add_tokens(['[IMG0]', '[IMG1]', '[IMG2]', '[IMG3]', '[IMG4]', '[IMG5]', '[IMG6]', '[IMG7]'], special_tokens=True)
+model.resize_token_embeddings(len(tokenizer))
+ckpt = T.load('_ckpt/mgie_7b/mllm.pt', map_location='cpu')
+model.load_state_dict(ckpt, strict=False)
+
+mm_use_im_start_end = getattr(model.config, 'mm_use_im_start_end', False)
+tokenizer.add_tokens([DEFAULT_IMAGE_PATCH_TOKEN], special_tokens=True)
+if mm_use_im_start_end: tokenizer.add_tokens([DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN], special_tokens=True)
+
+vision_tower = model.get_model().vision_tower[0]
+vision_tower = transformers.CLIPVisionModel.from_pretrained(vision_tower.config._name_or_path, torch_dtype=T.float16, low_cpu_mem_usage=True).cuda()
+model.get_model().vision_tower[0] = vision_tower
+vision_config = vision_tower.config
+vision_config.im_patch_token = tokenizer.convert_tokens_to_ids([DEFAULT_IMAGE_PATCH_TOKEN])[0]
+vision_config.use_im_start_end = mm_use_im_start_end
+if mm_use_im_start_end: vision_config.im_start_token, vision_config.im_end_token = tokenizer.convert_tokens_to_ids([DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN])
+image_token_len = (vision_config.image_size//vision_config.patch_size)**2
+
+_ = model.eval()
+
+pipe = diffusers.StableDiffusionInstructPix2PixPipeline.from_pretrained('timbrooks/instruct-pix2pix', torch_dtype=T.float16).to('cuda')
+pipe.set_progress_bar_config(disable=True)
+pipe.unet.load_state_dict(T.load('_ckpt/mgie_7b/unet.pt', map_location='cpu'))
+print('--init MGIE--')
+
+@spaces.GPU(enable_queue=True)
+def go_mgie(img, txt, seed, cfg_txt, cfg_img):
+    EMB = ckpt['emb'].cuda()
+    with T.inference_mode(): NULL = model.edit_head(T.zeros(1, 8, 4096).half().to('cuda'), EMB)
         
-        with gr.Tab("Generate Images"):
-            with gr.Row():
-                with gr.Column():
-                    prompt = gr.Textbox(label="Prompt")
-                    num_samples = gr.Number(label="Number of Samples", value=1)
-                    guidance_scale = gr.Number(label="Guidance Scale", value=7.5)
-                    height = gr.Number(label="Height", value=512)
-                    width = gr.Number(label="Width", value=512)
-                    num_inference_steps = gr.Slider(label="Steps", value=50, minimum=1, maximum=100)
-                    generate_button = gr.Button("Generate Images")
-                with gr.Column():
-                    gallery = gr.Gallery(label="Generated Images")
-                generate_button.click(generate_images, inputs=[prompt, num_samples, height, width, num_inference_steps, guidance_scale], outputs=gallery)
+    img, seed = crop_resize(Image.fromarray(img).convert('RGB')), int(seed)
+    inp = img
+
+    img = image_processor.preprocess(img, return_tensors='pt')['pixel_values'][0]
+    txt = "what will this image be like if '%s'"%(txt)
+    txt = txt+'\n'+DEFAULT_IM_START_TOKEN+DEFAULT_IMAGE_PATCH_TOKEN*image_token_len+DEFAULT_IM_END_TOKEN
+    conv = conv_templates['vicuna_v1_1'].copy()
+    conv.append_message(conv.roles[0], txt), conv.append_message(conv.roles[1], None)
+    txt = conv.get_prompt()
+    txt = tokenizer(txt)
+    txt, mask = T.as_tensor(txt['input_ids']), T.as_tensor(txt['attention_mask'])
+
+    with T.inference_mode():
+        _ = model.cuda()
+        out = model.generate(txt.unsqueeze(dim=0).cuda(), images=img.half().unsqueeze(dim=0).cuda(), attention_mask=mask.unsqueeze(dim=0).cuda(), 
+                             do_sample=False, max_new_tokens=96, num_beams=1, no_repeat_ngram_size=3, 
+                             return_dict_in_generate=True, output_hidden_states=True)
+        out, hid = out['sequences'][0].tolist(), T.cat([x[-1] for x in out['hidden_states']], dim=1)[0]
         
-    demo.launch()
+        if 32003 in out: p = out.index(32003)-1
+        else: p = len(hid)-9
+        p = min(p, len(hid)-9)
+        hid = hid[p:p+8]
+
+        out = remove_alter(tokenizer.decode(out))
+        _ = model.cuda()
+        emb = model.edit_head(hid.unsqueeze(dim=0), EMB)
+        res = pipe(image=inp, prompt_embeds=emb, negative_prompt_embeds=NULL, 
+                   generator=T.Generator(device='cuda').manual_seed(seed), guidance_scale=cfg_txt, image_guidance_scale=cfg_img).images[0]
 
-if __name__ == "__main__":
-    gradio_app()
+    return res, out
+
+go_mgie(np.array(Image.open('./_input/0.jpg').convert('RGB')), 'make the frame red', 13331, 7.5, 1.5)
+print('--init GO--')
+
+with gr.Blocks() as app:
+    gr.Markdown(
+        """
+        # MagiX: Edit Personalized Images using Gen AI
+        """
+    )
+    with gr.Row(): inp, res = [gr.Image(height=384, width=384, label='Input Image', interactive=True), 
+                               gr.Image(height=384, width=384, label='Goal Image', interactive=True)]
+    with gr.Row(): txt, out = [gr.Textbox(label='Instruction', interactive=True), 
+                               gr.Textbox(label='Expressive Instruction', interactive=False)]
+    with gr.Row(): seed, cfg_txt, cfg_img = [gr.Number(value=13331, label='Seed', interactive=True), 
+                                             gr.Number(value=7.5, label='Text CFG', interactive=True), 
+                                             gr.Number(value=1.5, label='Image CFG', interactive=True)]
+    with gr.Row(): btn_sub = gr.Button('Submit')
+    btn_sub.click(fn=go_mgie, inputs=[inp, txt, seed, cfg_txt, cfg_img], outputs=[res, out])
+    
+app.launch()