Update app.py

app.py

@@ -96,15 +96,15 @@ isomer_color_weights = torch.from_numpy(np.array([1, 0.5, 1, 0.5])).float().to(d
 
 # model initialization and loading
 # flux
-taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=torch.bfloat16).to(device_0)
-good_vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae", torch_dtype=torch.bfloat16, token=access_token).to(device_0)
-# flux_pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16, token=access_token).to(device=device_0, dtype=torch.bfloat16)
-flux_pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16, vae=taef1, token=access_token).to(device_0)
+# taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=torch.bfloat16).to(device_0)
+# good_vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae", torch_dtype=torch.bfloat16, token=access_token).to(device_0)
+flux_pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16, token=access_token).to(device=device_0, dtype=torch.bfloat16)
+# flux_pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16, vae=taef1, token=access_token).to(device_0)
 flux_lora_ckpt_path = hf_hub_download(repo_id="LTT/xxx-ckpt", filename="rgb_normal_large.safetensors", repo_type="model", token=access_token)
 flux_pipe.load_lora_weights(flux_lora_ckpt_path)
 flux_pipe.to(device=device_0, dtype=torch.bfloat16)
 torch.cuda.empty_cache()
-flux_pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(flux_pipe)
+# flux_pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(flux_pipe)
 
 
 # lrm
@@ -208,28 +208,28 @@ def generate_multi_view_images(prompt, seed):
     # generator = torch.manual_seed(seed)
     generator = torch.Generator().manual_seed(seed)
     with torch.no_grad():
-        # images = flux_pipe(
-        #     prompt=prompt,
-        #     num_inference_steps=10,
-        #     guidance_scale=3.5,
-        #     num_images_per_prompt=1,
-        #     width=resolution * 4,
-        #     height=resolution * 2,
-        #     output_type='np',
-        #     generator=generator,
-        #     good_vae=good_vae,
-        # ).images
-        for img in flux_pipe.flux_pipe_call_that_returns_an_iterable_of_images(
+        images = flux_pipe(
             prompt=prompt,
+            num_inference_steps=5,
             guidance_scale=3.5,
-            num_inference_steps=4,
+            num_images_per_prompt=1,
             width=resolution * 4,
             height=resolution * 2,
+            output_type='np',
             generator=generator,
-            output_type="np",
             good_vae=good_vae,
-        ):
-            pass
+        ).images
+        # for img in flux_pipe.flux_pipe_call_that_returns_an_iterable_of_images(
+        #     prompt=prompt,
+        #     guidance_scale=3.5,
+        #     num_inference_steps=4,
+        #     width=resolution * 4,
+        #     height=resolution * 2,
+        #     generator=generator,
+        #     output_type="np",
+        #     good_vae=good_vae,
+        # ):
+        #     pass
     # 返回最终的图像和种子（通过外部调用处理）
     return img