Update app.py

app.py

@@ -6,7 +6,6 @@ import gradio as gr
 import torch
 from einops import rearrange
 from PIL import Image
-from transformers import pipeline
 
 from flux.cli import SamplingOptions
 from flux.sampling import denoise, get_noise, get_schedule, prepare, unpack
@@ -16,6 +15,7 @@ from pulid.utils import resize_numpy_image_long
 
 NSFW_THRESHOLD = 0.85
 
+
 def get_models(name: str, device: torch.device, offload: bool):
     t5 = load_t5(device, max_length=128)
     clip = load_clip(device)
@@ -27,15 +27,17 @@ def get_models(name: str, device: torch.device, offload: bool):
 
 class FluxGenerator:
     def __init__(self):
-        self.device = torch.device('cuda')
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         self.offload = False
-        self.model_name = 'flux-dev'
+        self.model_name = "flux-dev"
         self.model, self.ae, self.t5, self.clip = get_models(
             self.model_name,
             device=self.device,
             offload=self.offload,
         )
-        self.pulid_model = PuLIDPipeline(self.model, 'cuda', weight_dtype=torch.bfloat16)
+        device_str = "cuda" if torch.cuda.is_available() else "cpu"
+        weight_dtype = torch.bfloat16 if device_str == "cuda" else torch.float32
+        self.pulid_model = PuLIDPipeline(self.model, device_str, weight_dtype=weight_dtype)
         self.pulid_model.load_pretrain()
 
 
@@ -45,19 +47,19 @@ flux_generator = FluxGenerator()
 @spaces.GPU
 @torch.inference_mode()
 def generate_image(
-        width,
-        height,
-        num_steps,
-        start_step,
-        guidance,
-        seed,
-        prompt,
-        id_image=None,
-        id_weight=1.0,
-        neg_prompt="",
-        true_cfg=1.0,
-        timestep_to_start_cfg=1,
-        max_sequence_length=128,
+    width,
+    height,
+    num_steps,
+    start_step,
+    guidance,
+    seed,
+    prompt,
+    id_image=None,
+    id_weight=1.0,
+    neg_prompt="",
+    true_cfg=1.0,
+    timestep_to_start_cfg=1,
+    max_sequence_length=128,
 ):
     flux_generator.t5.max_length = max_sequence_length
 
@@ -83,7 +85,9 @@ def generate_image(
 
     if id_image is not None:
         id_image = resize_numpy_image_long(id_image, 1024)
-        id_embeddings, uncond_id_embeddings = flux_generator.pulid_model.get_id_embedding(id_image, cal_uncond=use_true_cfg)
+        id_embeddings, uncond_id_embeddings = flux_generator.pulid_model.get_id_embedding(
+            id_image, cal_uncond=use_true_cfg
+        )
     else:
         id_embeddings = None
         uncond_id_embeddings = None
@@ -96,7 +100,7 @@ def generate_image(
         opts.height,
         opts.width,
         device=flux_generator.device,
-        dtype=torch.bfloat16,
+        dtype=torch.bfloat16 if flux_generator.device.type == "cuda" else torch.float32,
         seed=opts.seed,
     )
     print(x)
@@ -107,7 +111,10 @@ def generate_image(
     )
 
     if flux_generator.offload:
-        flux_generator.t5, flux_generator.clip = flux_generator.t5.to(flux_generator.device), flux_generator.clip.to(flux_generator.device)
+        flux_generator.t5, flux_generator.clip = (
+            flux_generator.t5.to(flux_generator.device),
+            flux_generator.clip.to(flux_generator.device),
+        )
     inp = prepare(t5=flux_generator.t5, clip=flux_generator.clip, img=x, prompt=opts.prompt)
     inp_neg = prepare(t5=flux_generator.t5, clip=flux_generator.clip, img=x, prompt=neg_prompt) if use_true_cfg else None
 
@@ -119,8 +126,15 @@ def generate_image(
 
     # denoise initial noise
     x = denoise(
-        flux_generator.model, **inp, timesteps=timesteps, guidance=opts.guidance, id=id_embeddings, id_weight=id_weight,
-        start_step=start_step, uncond_id=uncond_id_embeddings, true_cfg=true_cfg,
+        flux_generator.model,
+        **inp,
+        timesteps=timesteps,
+        guidance=opts.guidance,
+        id=id_embeddings,
+        id_weight=id_weight,
+        start_step=start_step,
+        uncond_id=uncond_id_embeddings,
+        true_cfg=true_cfg,
         timestep_to_start_cfg=timestep_to_start_cfg,
         neg_txt=inp_neg["txt"] if use_true_cfg else None,
         neg_txt_ids=inp_neg["txt_ids"] if use_true_cfg else None,
@@ -135,7 +149,10 @@ def generate_image(
 
     # decode latents to pixel space
     x = unpack(x.float(), opts.height, opts.width)
-    with torch.autocast(device_type=flux_generator.device.type, dtype=torch.bfloat16):
+    with torch.autocast(
+        device_type=flux_generator.device.type,
+        dtype=torch.bfloat16 if flux_generator.device.type == "cuda" else torch.float32,
+    ):
         x = flux_generator.ae.decode(x)
 
     if flux_generator.offload:
@@ -147,15 +164,13 @@ def generate_image(
     print(f"Done in {t1 - t0:.1f}s.")
     # bring into PIL format
     x = x.clamp(-1, 1)
-    # x = embed_watermark(x.float())
     x = rearrange(x[0], "c h w -> h w c")
 
     img = Image.fromarray((127.5 * (x + 1.0)).cpu().byte().numpy())
     return img, str(opts.seed), flux_generator.pulid_model.debug_img_list
 
-def create_demo(args, model_name: str, device: str = "cuda" if torch.cuda.is_available() else "cpu",
-                offload: bool = False):
-    
+
+def create_demo(args, model_name: str, device: str = "cuda" if torch.cuda.is_available() else "cpu", offload: bool = False):
     with gr.Blocks(theme="soft") as demo:
         gr.HTML(
             """
@@ -163,14 +178,14 @@ def create_demo(args, model_name: str, device: str = "cuda" if torch.cuda.is_ava
                 <a href="https://huggingface.co/spaces/openfree/Best-AI" target="_blank">
                     <img src="https://img.shields.io/static/v1?label=OpenFree&message=BEST%20AI%20Services&color=%230000ff&labelColor=%23000080&logo=huggingface&logoColor=%23ffa500&style=for-the-badge" alt="OpenFree badge">
                 </a>
-        
+
                 <a href="https://discord.gg/openfreeai" target="_blank">
                     <img src="https://img.shields.io/static/v1?label=Discord&message=Openfree%20AI&color=%230000ff&labelColor=%23800080&logo=discord&logoColor=white&style=for-the-badge" alt="Discord badge">
                 </a>
             </div>
             """
         )
-        
+
         with gr.Row():
             with gr.Column():
                 prompt = gr.Textbox(label="Prompt", value="portrait, color, cinematic")
@@ -183,75 +198,102 @@ def create_demo(args, model_name: str, device: str = "cuda" if torch.cuda.is_ava
                 start_step = gr.Slider(0, 10, 0, step=1, label="timestep to start inserting ID")
                 guidance = gr.Slider(1.0, 10.0, 4, step=0.1, label="Guidance")
                 seed = gr.Textbox(-1, label="Seed (-1 for random)")
-                max_sequence_length = gr.Slider(128, 512, 128, step=128,
-                                                label="max_sequence_length for prompt (T5), small will be faster")
+                max_sequence_length = gr.Slider(128, 512, 128, step=128, label="max_sequence_length for prompt (T5), small will be faster")
 
-                with gr.Accordion("Advanced Options (True CFG, true_cfg_scale=1 means use fake CFG, >1 means use true CFG, if using true CFG, we recommend set the guidance scale to 1)", open=False):    # noqa E501
+                with gr.Accordion(
+                    "Advanced Options (True CFG, true_cfg_scale=1 means use fake CFG, >1 means use true CFG, if using true CFG, we recommend set the guidance scale to 1)",
+                    open=False,
+                ):
                     neg_prompt = gr.Textbox(
                         label="Negative Prompt",
-                        value="bad quality, worst quality, text, signature, watermark, extra limbs")
+                        value="bad quality, worst quality, text, signature, watermark, extra limbs",
+                    )
                     true_cfg = gr.Slider(1.0, 10.0, 1, step=0.1, label="true CFG scale")
                     timestep_to_start_cfg = gr.Slider(0, 20, 1, step=1, label="timestep to start cfg", visible=args.dev)
 
                 generate_btn = gr.Button("Generate")
-                
+
             with gr.Column():
                 output_image = gr.Image(label="Generated Image")
                 seed_output = gr.Textbox(label="Used Seed")
-                intermediate_output = gr.Gallery(label='Output', elem_id="gallery", visible=args.dev)
+                intermediate_output = gr.Gallery(label="Output", elem_id="gallery", visible=args.dev)
 
         with gr.Row(), gr.Column():
             gr.Markdown("## Examples")
             example_inps = [
                 [
-                    'a woman holding sign with glowing green text \"PuLID for FLUX\"',
-                    'example_inputs/qw1.webp',
-                    4, 4, 2680261499100305976, 1
+                    'a woman holding sign with glowing green text "PuLID for FLUX"',
+                    "example_inputs/qw1.webp",
+                    4,
+                    4,
+                    2680261499100305976,
+                    1,
                 ],
                 [
-                    'portrait, pixar',
-                    'example_inputs/qw2.webp',
-                    1, 4, 9445036702517583939, 1
+                    "portrait, pixar",
+                    "example_inputs/qw2.webp",
+                    1,
+                    4,
+                    9445036702517583939,
+                    1,
                 ],
             ]
-            gr.Examples(examples=example_inps, inputs=[prompt, id_image, start_step, guidance, seed, true_cfg],
-                        label='fake CFG')
+            gr.Examples(examples=example_inps, inputs=[prompt, id_image, start_step, guidance, seed, true_cfg], label="fake CFG")
 
             example_inps = [
                 [
-                    'portrait, made of ice sculpture',
-                    'example_inputs/qw3.webp',
-                    1, 1, 3811899118709451814, 5
+                    "portrait, made of ice sculpture",
+                    "example_inputs/qw3.webp",
+                    1,
+                    1,
+                    3811899118709451814,
+                    5,
                 ],
             ]
-            gr.Examples(examples=example_inps, inputs=[prompt, id_image, start_step, guidance, seed, true_cfg],
-                        label='true CFG')
+            gr.Examples(examples=example_inps, inputs=[prompt, id_image, start_step, guidance, seed, true_cfg], label="true CFG")
 
         generate_btn.click(
             fn=generate_image,
-            inputs=[width, height, num_steps, start_step, guidance, seed, prompt, id_image, id_weight, neg_prompt,
-                    true_cfg, timestep_to_start_cfg, max_sequence_length],
+            inputs=[
+                width,
+                height,
+                num_steps,
+                start_step,
+                guidance,
+                seed,
+                prompt,
+                id_image,
+                id_weight,
+                neg_prompt,
+                true_cfg,
+                timestep_to_start_cfg,
+                max_sequence_length,
+            ],
             outputs=[output_image, seed_output, intermediate_output],
         )
 
     return demo
 
+
 if __name__ == "__main__":
     import argparse
 
     parser = argparse.ArgumentParser(description="PuLID for FLUX.1-dev")
-    parser.add_argument("--name", type=str, default="flux-dev", choices=list('flux-dev'),
-                        help="currently only support flux-dev")
-    parser.add_argument("--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu",
-                        help="Device to use")
+    parser.add_argument("--name", type=str, default="flux-dev", choices=["flux-dev"], help="currently only support flux-dev")
+    parser.add_argument(
+        "--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu", help="Device to use"
+    )
     parser.add_argument("--offload", action="store_true", help="Offload model to CPU when not in use")
     parser.add_argument("--port", type=int, default=8080, help="Port to use")
-    parser.add_argument("--dev", action='store_true', help="Development mode")
-    parser.add_argument("--pretrained_model", type=str, help='for development')
+    parser.add_argument("--dev", action="store_true", help="Development mode")
+    parser.add_argument("--pretrained_model", type=str, help="for development")
     args = parser.parse_args()
 
     import huggingface_hub
-    huggingface_hub.login(os.getenv('HF_TOKEN'))
+
+    hf_token = os.getenv("HF_TOKEN")
+    if hf_token:
+        huggingface_hub.login(hf_token)
 
     demo = create_demo(args, args.name, args.device, args.offload)
-    demo.launch()
+    demo.launch()