<feat> complete app.py.

.gitignore

@@ -0,0 +1 @@
+__pycache__/

app.py

@@ -18,10 +18,13 @@ from peft import LoraConfig
 from omegaconf import OmegaConf
 from safetensors.torch import safe_open
 from PIL import Image, ImageDraw, ImageFilter
+from huggingface_hub import hf_hub_download
+from transformers import pipeline
 
 from models import HunyuanVideoTransformer3DModel
 from pipelines import HunyuanVideoImageToVideoPipeline
 
+
 header = """
 # DRA-Ctrl Gradio App
 
@@ -33,9 +36,267 @@ header = """
 </div>
 """
 
+notice = """
+For easier testing, in spatially-aligned image generation tasks, when passing the condition image to `gradio_app`, 
+there's no need to manually input edge maps, depth maps, or other condition images - only the original image is required. 
+The corresponding condition images will be automatically extracted.
+"""
+
+
+
+@spaces.GPU
+def process_image_and_text(condition_image, target_prompt, condition_image_prompt, task):
+    # init models
+    transformer = HunyuanVideoTransformer3DModel.from_pretrained('hunyuanvideo-community/HunyuanVideo-I2V', 
+                                                                 subfolder="transformer", 
+                                                                 inference_subject_driven=task in ['subject_driven'])
+    scheduler = diffusers.FlowMatchEulerDiscreteScheduler()
+    vae = diffusers.AutoencoderKLHunyuanVideo.from_pretrained('hunyuanvideo-community/HunyuanVideo-I2V', 
+                                                              subfolder="vae")
+    text_encoder = transformers.LlavaForConditionalGeneration.from_pretrained('hunyuanvideo-community/HunyuanVideo-I2V', 
+                                                                              subfolder="text_encoder")
+    text_encoder_2 = transformers.CLIPTextModel.from_pretrained('hunyuanvideo-community/HunyuanVideo-I2V', 
+                                                                subfolder="text_encoder_2")
+    tokenizer = transformers.AutoTokenizer.from_pretrained('hunyuanvideo-community/HunyuanVideo-I2V', 
+                                                           subfolder="tokenizer")
+    tokenizer_2 = transformers.CLIPTokenizer.from_pretrained('hunyuanvideo-community/HunyuanVideo-I2V', 
+                                                             subfolder="tokenizer_2")
+    image_processor = transformers.CLIPImageProcessor.from_pretrained('hunyuanvideo-community/HunyuanVideo-I2V', 
+                                                                      subfolder="image_processor")
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    weight_dtype = torch.bfloat16
+
+    transformer.requires_grad_(False)
+    vae.requires_grad_(False).to(device, dtype=weight_dtype)
+    text_encoder.requires_grad_(False).to(device, dtype=weight_dtype)
+    text_encoder_2.requires_grad_(False).to(device, dtype=weight_dtype)
+    transformer.to(device, dtype=weight_dtype)
+    vae.enable_tiling()
+    vae.enable_slicing()
+
+    # insert LoRA
+    lora_config = LoraConfig(
+        r=16,
+        lora_alpha=16,
+        init_lora_weights="gaussian",
+        target_modules=[
+            'attn.to_k', 'attn.to_q', 'attn.to_v', 'attn.to_out.0',
+            'attn.add_k_proj', 'attn.add_q_proj', 'attn.add_v_proj', 'attn.to_add_out',
+            'ff.net.0.proj', 'ff.net.2',
+            'ff_context.net.0.proj', 'ff_context.net.2',
+            'norm1_context.linear', 'norm1.linear',
+            'norm.linear', 'proj_mlp', 'proj_out',
+        ]
+    )
+    transformer.add_adapter(lora_config)
+
+    # hack LoRA forward
+    def create_hacked_forward(module):
+        lora_forward = module.forward
+        non_lora_forward = module.base_layer.forward
+        img_sequence_length = int((args.img_size / 8 / 2) ** 2)
+        encoder_sequence_length = 144 + 252 # encoder sequence: 144 img 252 txt
+        num_imgs = 4
+        num_generated_imgs = 3
+        num_encoder_sequences = 2 if args.task in ['subject_driven', 'style_transfer'] else 1
+
+        def hacked_lora_forward(self, x, *args, **kwargs):
+            if x.shape[1] == img_sequence_length * num_imgs and len(x.shape) > 2:
+                return torch.cat((
+                    lora_forward(x[:, :-img_sequence_length*num_generated_imgs], *args, **kwargs),
+                    non_lora_forward(x[:, -img_sequence_length*num_generated_imgs:], *args, **kwargs)
+                ), dim=1)
+            elif x.shape[1] == encoder_sequence_length * num_encoder_sequences or x.shape[1] == encoder_sequence_length:
+                return lora_forward(x, *args, **kwargs)
+            elif x.shape[1] == img_sequence_length * num_imgs + encoder_sequence_length * num_encoder_sequences:
+                return torch.cat((
+                    lora_forward(x[:, :(num_imgs - num_generated_imgs)*img_sequence_length], *args, **kwargs),
+                    non_lora_forward(x[:, (num_imgs - num_generated_imgs)*img_sequence_length:-num_encoder_sequences*encoder_sequence_length], *args, **kwargs),
+                    lora_forward(x[:, -num_encoder_sequences*encoder_sequence_length:], *args, **kwargs)
+                ), dim=1)
+            elif x.shape[1] == 3072:
+                return non_lora_forward(x, *args, **kwargs)
+            else:
+                raise ValueError(
+                    f"hacked_lora_forward receives unexpected sequence length: {x.shape[1]}, input shape: {x.shape}!"
+                )
+
+        return hacked_lora_forward.__get__(module, type(module))
+
+    for n, m in transformer.named_modules():
+        if isinstance(m, peft.tuners.lora.layer.Linear):
+            m.forward = create_hacked_forward(m)
+
+    # load LoRA weights
+    model_root = hf_hub_download(
+        repo_id="Kunbyte/DRA-Ctrl",
+        filename=f"{task}.safetensors",
+        resume_download=True)
+
+    try:
+        with safe_open(model_root, framework="pt") as f:
+            lora_weights = {}
+            for k in f.keys():
+                param = f.get_tensor(k) 
+                if k.endswith(".weight"):
+                    k = k.replace('.weight', '.default.weight')
+                lora_weights[k] = param
+            transformer.load_state_dict(lora_weights, strict=False)
+    except Exception as e:
+        raise ValueError(f'{e}')
+
+    transformer.requires_grad_(False)
+
+    pipe = HunyuanVideoImageToVideoPipeline(
+        text_encoder=text_encoder,
+        tokenizer=tokenizer,
+        transformer=transformer,
+        vae=vae,
+        scheduler=copy.deepcopy(scheduler),
+        text_encoder_2=text_encoder_2,
+        tokenizer_2=tokenizer_2,
+        image_processor=image_processor,
+    )
+
+    # start generation
+    c_txt = None if condition_image_prompt == "" else condition_image_prompt
+    c_img = condition_image.resize((512, 512))
+    t_txt = target_prompt
+
+    if args.task not in ['subject_driven', 'style_transfer']:
+        if args.task == "canny":
+            def get_canny_edge(img):
+                img_np = np.array(img)
+                img_gray = cv2.cvtColor(img_np, cv2.COLOR_BGR2GRAY)
+                edges = cv2.Canny(img_gray, 100, 200)
+                edges_tmp = Image.fromarray(edges).convert("RGB")
+                edges_tmp.save(os.path.join(save_dir, f"edges.png"))
+                edges[edges == 0] = 128
+                return Image.fromarray(edges).convert("RGB")
+            c_img = get_canny_edge(c_img)
+        elif args.task == "coloring":
+            c_img = (
+                c_img.resize((args.img_size, args.img_size))
+                .convert("L")
+                .convert("RGB")
+            )
+        elif args.task == "deblurring":
+            blur_radius = 10
+            c_img = (
+                c_img.convert("RGB")
+                .filter(ImageFilter.GaussianBlur(blur_radius))
+                .resize((args.img_size, args.img_size))
+                .convert("RGB")
+            )
+        elif args.task == "depth":
+            def get_depth_map(img):
+                from transformers import pipeline
+
+                depth_pipe = pipeline(
+                    task="depth-estimation",
+                    model="LiheYoung/depth-anything-small-hf",
+                    device="cpu",
+                )
+                return depth_pipe(img)["depth"].convert("RGB").resize((args.img_size, args.img_size))
+            c_img = get_depth_map(c_img)
+            c_img.save(os.path.join(save_dir, f"depth.png"))
+            k = (255 - 128) / 255
+            b = 128
+            c_img = c_img.point(lambda x: k * x + b)
+        elif args.task == "depth_pred":
+            c_img = c_img
+        elif args.task == "fill":
+            c_img = c_img.resize((args.img_size, args.img_size)).convert("RGB")
+            x1, x2 = args.fill_x1, args.fill_x2
+            y1, y2 = args.fill_y1, args.fill_y2
+            mask = Image.new("L", (args.img_size, args.img_size), 0)
+            draw = ImageDraw.Draw(mask)
+            draw.rectangle((x1, y1, x2, y2), fill=255)
+            if args.inpainting:
+                mask = Image.eval(mask, lambda a: 255 - a)
+            c_img = Image.composite(
+                c_img,
+                Image.new("RGB", (args.img_size, args.img_size), (255, 255, 255)),
+                mask
+            )
+            c_img.save(os.path.join(save_dir, f"mask.png"))
+            c_img = Image.composite(
+                c_img,
+                Image.new("RGB", (args.img_size, args.img_size), (128, 128, 128)),
+                mask
+            )
+        elif args.task == "sr":
+            c_img = c_img.resize((int(args.img_size / 4), int(args.img_size / 4))).convert("RGB")
+            c_img.save(os.path.join(save_dir, f"low_resolution.png"))
+            c_img = c_img.resize((args.img_size, args.img_size))
+            c_img.save(os.path.join(save_dir, f"low_to_high.png"))
+
+    gen_img = pipe(
+        image=c_img,
+        prompt=[t_txt.strip()],
+        prompt_condition=[c_txt.strip()] if c_txt is not None else None,
+        prompt_2=[t_txt],
+        height=512,
+        width=512,
+        num_frames=5,
+        num_inference_steps=50,
+        guidance_scale=6.0,
+        num_videos_per_prompt=1,
+        generator=torch.Generator(device=pipe.transformer.device).manual_seed(0),
+        output_type='pt',
+        image_embed_interleave=4,
+        frame_gap=48,
+        mixup=True,
+        mixup_num_imgs=2,
+    ).frames
+
+    gen_img = gen_img[:, 0:1, :, :, :]
+    gen_img = gen_img.squeeze(0).squeeze(0).cpu().to(torch.float32).numpy()
+    gen_img = np.transpose(gen_img, (1, 2, 0))
+    gen_img = (gen_img * 255).astype(np.uint8)
+    gen_img = Image.fromarray(gen_img)
+
+    return gen_img
+
 def create_app():
     with gr.Blocks() as app:
         gr.Markdown(header, elem_id="header")
+        with gr.Row(equal_height=False):
+            with gr.Column(variant="panel", elem_classes="inputPanel"):
+                condition_image = gr.Image(
+                    type="pil", label="Condition Image", width=300, elem_id="input"
+                )
+                task = gr.Radio(
+                    [
+                        ("Subject-driven Image Generation", "subject_driven"),
+                        ("Canny-to-Image", "canny"),
+                        ("Colorization", "coloring"),
+                        ("Deblurring", "deblurring"),
+                        ("Depth-to-Image", "depth"),
+                        ("Depth Prediction", "depth_pred"),
+                        ("In/Out-Painting", "fill"),
+                        ("Super-Resolution", "sr"),
+                        ("Style Transfer", "style_transfer")
+                    ],
+                    label="Task Selection",
+                    value="subject_driven",
+                    interactive=True,
+                    elem_id="task_selection"
+                )
+                gr.Markdown(notice, elem_id="notice")
+                target_prompt = gr.Textbox(lines=2, label="Target Prompt", elem_id="text")
+                condition_image_prompt = gr.Textbox(lines=2, label="Condition Image Prompt", elem_id="text")
+                submit_btn = gr.Button("Run", elem_id="submit_btn")
+
+            with gr.Column(variant="panel", elem_classes="outputPanel"):
+                output_image = gr.Image(type="pil", elem_id="output")
+
+        submit_btn.click(
+            fn=process_image_and_text,
+            inputs=[condition_image, target_prompt, condition_image_prompt, task],
+            outputs=output_image,
+        )
         
     return app