22

app.py

@@ -18,16 +18,26 @@ from utils.prompt_utils import polish_prompt
 #     subprocess.run(["git", "clone", GIT_URL])
 # else:
 #     print("Folder already exists.")
-
+MODEL_REPO = "meituan-longcat/LongCat-Image"
 MAX_SEED = np.iinfo(np.int32).max
 
-pipe = LongCatImagePipeline.from_pretrained(
-    "meituan-longcat/LongCat-Image",
-    torch_dtype=torch.bfloat16,
-    low_cpu_mem_usage=False,
+text_processor = AutoProcessor.from_pretrained( 
+    MODEL_REPO, 
+    subfolder = 'tokenizer'
 )
-pipe.to("cuda")
+transformer = LongCatImageTransformer2DModel.from_pretrained( 
+    MODEL_REPO , 
+    subfolder = 'transformer', 
+    torch_dtype=torch.bfloat16, 
+    use_safetensors=True
+).to("cuda")
 
+pipe = LongCatImagePipeline.from_pretrained(
+    MODEL_REPO,
+    transformer=transformer,
+    text_processor=text_processor
+)
+pipe.to("cuda", torch.bfloat16)
 
 def prepare(prompt, is_polish_prompt):
     if not is_polish_prompt: return prompt, False