Update app.py

app.py

@@ -30,9 +30,9 @@ transformer_trainable_parameters = None
 
 def load_lora_from_subfolder():
     repo_id = "benzweijia/Adv-GRPO"
-    subfolder = "PickScore"          
+    subfolder = "DINO"          
 
-    local_dir = "/tmp/PickScore"
+    local_dir = "/tmp/DINO"
     os.makedirs(local_dir, exist_ok=True)
 
     for filename in ["adapter_config.json", "adapter_model.safetensors"]:
@@ -107,7 +107,7 @@ def init_model():
     pipeline.text_encoder_2.to("cuda")
     pipeline.text_encoder_3.to("cuda")
     pipeline.transformer.to("cuda")
-    config.train.lora_path = "benzweijia/Adv-GRPO/PickScore"
+    config.train.lora_path = "benzweijia/Adv-GRPO/DINO"
     config.use_lora = True
     lora_dir = load_lora_from_subfolder()
 
@@ -115,7 +115,7 @@ def init_model():
         print("🔥 Loading LoRA from:", config.train.lora_path)
         pipeline.transformer = PeftModel.from_pretrained(
             pipeline.transformer,
-            os.path.join(lora_dir,"PickScore")
+            os.path.join(lora_dir,"DINO")
         )
         pipeline.transformer.set_adapter("default")
 
@@ -146,7 +146,6 @@ def infer(prompt):
     if pipeline is None:
         init_model()
     print(pipeline)
-    print("start infer 1111")
         
 
     prompts = [prompt]
@@ -157,7 +156,6 @@ def infer(prompt):
         max_sequence_length=128,
         device="cuda"
     )
-    print("start infer 2")
 
     neg_embed, neg_pooled_embed = compute_text_embeddings(
         [""], text_encoders, tokenizers,
@@ -167,7 +165,6 @@ def infer(prompt):
 
     neg_prompt_embeds = neg_embed.repeat(1, 1, 1)
     neg_pooled_prompt_embeds = neg_pooled_embed.repeat(1, 1)
-    print("start infer 3")
 
     # generation seed
     generator = torch.Generator().manual_seed(0)
@@ -192,9 +189,6 @@ def infer(prompt):
             generator=generator,
         )
         
-    print("images type:", type(images))
-    print("images len:", len(images))
-    print("first image shape:", images[0].shape)
 
     # Convert to PIL
     pil = Image.fromarray(
@@ -216,7 +210,7 @@ demo = gr.Interface(
     fn=infer,
     inputs=gr.Textbox(lines=2, label="Prompt"),
     outputs=gr.Image(type="pil"),
-    title="Adv-GRPO(PickScore)",
+    title="Adv-GRPO(DINO)",
     description="Enter a prompt and generate image using Adv-GRPO",
 )