Add project files

scripts/inference_edit.py

@@ -0,0 +1,37 @@
+import torch
+from PIL import Image
+from transformers import AutoProcessor
+from longcat_image.models import LongCatImageTransformer2DModel
+from longcat_image.pipelines import LongCatImageEditPipeline
+
+if __name__ == '__main__':
+
+    device = torch.device('cuda')
+    checkpoint_dir = './weights/LongCat-Image-Edit'
+    text_processor = AutoProcessor.from_pretrained( checkpoint_dir, subfolder = 'tokenizer'  )
+    transformer = LongCatImageTransformer2DModel.from_pretrained( checkpoint_dir , subfolder = 'transformer', 
+        torch_dtype=torch.bfloat16, use_safetensors=True).to(device)
+
+    pipe = LongCatImageEditPipeline.from_pretrained(
+        checkpoint_dir,
+        transformer=transformer,
+        text_processor=text_processor,
+        torch_dtype=torch.bfloat16,
+    )
+    # pipe.to(device, torch.bfloat16)  # Uncomment for high VRAM devices (Faster inference)
+    pipe.enable_model_cpu_offload()  # Offload to CPU to save VRAM (Required ~19 GB); slower but prevents OOM
+
+    generator = torch.Generator("cpu").manual_seed(43)
+    img = Image.open('assets/test.png').convert('RGB')
+    prompt = '将猫变成狗'
+    image = pipe(
+        img,
+        prompt,
+        negative_prompt='',
+        guidance_scale=4.5,
+        num_inference_steps=50,
+        num_images_per_prompt=1,
+        generator=generator
+    ).images[0]
+
+    image.save('./edit_example.png')

scripts/inference_t2i.py

@@ -0,0 +1,44 @@
+import torch
+
+from transformers import AutoProcessor
+from longcat_image.models import LongCatImageTransformer2DModel
+from longcat_image.pipelines import LongCatImagePipeline
+from misc.prompt_rewrite_api import prompt_rewrite_deepseek
+
+if __name__ == '__main__':
+
+    device = torch.device('cuda')
+    checkpoint_dir = './weights/LongCat-Image'
+
+    text_processor = AutoProcessor.from_pretrained( checkpoint_dir, subfolder = 'tokenizer'  )
+    transformer = LongCatImageTransformer2DModel.from_pretrained( checkpoint_dir , subfolder = 'transformer', 
+        torch_dtype=torch.bfloat16, use_safetensors=True).to(device)
+
+    pipe = LongCatImagePipeline.from_pretrained(
+        checkpoint_dir,
+        transformer=transformer,
+        text_processor=text_processor,
+        torch_dtype=torch.bfloat16
+    )
+    # pipe.to(device, torch.bfloat16)  # Uncomment for high VRAM devices (Faster inference)
+    pipe.enable_model_cpu_offload()  # Offload to CPU to save VRAM (Required ~17 GB); slower but prevents OOM
+
+    prompt = '一个年轻的亚裔女性，身穿黄色针织衫，搭配白色项链。她的双手放在膝盖上，表情恬静。背景是一堵粗糙的砖墙，午后的阳光温暖地洒在她身上，营造出一种宁静而温馨的氛围。镜头采用中距离视角，突出她的神态和服饰的细节。光线柔和地打在她的脸上，强调她的五官和饰品的质感，增加画面的层次感与亲和力。整个画面构图简洁，砖墙的纹理与阳光的光影效果相得益彰，突显出人物的优雅与从容。'
+
+    enable_prompt_rewrite_api = False
+    if enable_prompt_rewrite_api:
+        prompt = prompt_rewrite_deepseek(prompt)
+
+    image = pipe(
+        prompt,
+        negative_prompt='',
+        height=768,
+        width=1344,
+        guidance_scale=4.5,
+        num_inference_steps=50,
+        num_images_per_prompt=1,
+        generator= torch.Generator("cpu").manual_seed(43),
+        enable_cfg_renorm=True,
+        enable_prompt_rewrite=True if not enable_prompt_rewrite_api else False
+    ).images[0]
+    image.save('./t2i_example.png')