Update gen_sd3.py

gen_sd3.py

@@ -22,19 +22,26 @@ json_filename = args.json_filename
 cuda_device = f"cuda:{args.cuda}"
 print(json_filename, cuda_device)
 
-image_dir = "/mnt/petrelfs/zhuchenglin/LLaVA/playground/data/LLaVA-Pretrain/images"
+image_dir = "/mnt/petrelfs/zhuchenglin/LLaVA/playground/data"
 with open(json_filename, "r") as f:
     json_data = json.load(f)
     
 pipe = StableDiffusion3Pipeline.from_pretrained("stabilityai/stable-diffusion-3-medium-diffusers", torch_dtype=torch.float16)
-pipe.to('cuda')
+pipe.to(cuda_device)
 
 for text in json_data:
+    prompt = ""
+    for caption in text['conversations']:
+        if caption['from'] == 'gpt':
+            prompt += caption['value']
+    # for caption in text['conversations']:
+    #     prompt += caption['value']
+    
     image = pipe(
-        prompt=text["conversations"][1]["value"],
-        prompt_3=text["conversations"][1]["value"],
+        prompt=prompt,
+        prompt_3=prompt,
         negative_prompt="",
-        num_inference_steps=100,
+        num_inference_steps=60,
         height=1024,
         width=1024,
         guidance_scale=10.0,