update

app.py

@@ -6,24 +6,90 @@ import torch
 import devicetorch
 from diffusers import DiffusionPipeline
 
-device = devicetorch.get(torch)
-if device == "cuda":
-    dtype = torch.bfloat16
-elif device == "mps":
-    dtype = torch.float16
-else:
-    dtype = torch.float32
-#dtype = torch.bfloat16
-#device = "cuda" if torch.cuda.is_available() else "cpu"
 
-#pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16, revision="refs/pr/1").to(device)
-pipe = DiffusionPipeline.from_pretrained("cocktailpeanut/xulf-s", torch_dtype=dtype).to(device)
+# Quant
+from optimum.quanto import freeze, qfloat8, quantize
+from diffusers import FlowMatchEulerDiscreteScheduler, AutoencoderKL
+from diffusers.models.transformers.transformer_flux import FluxTransformer2DModel
+from diffusers.pipelines.flux.pipeline_flux import FluxPipeline
+from transformers import CLIPTextModel, CLIPTokenizer,T5EncoderModel, T5TokenizerFast
+
+
+#device = devicetorch.get(torch)
+#if device == "cuda":
+#    dtype = torch.bfloat16
+#elif device == "mps":
+#    dtype = torch.float16
+#else:
+#    dtype = torch.float32
+##dtype = torch.bfloat16
+##device = "cuda" if torch.cuda.is_available() else "cpu"
+#
+##pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16, revision="refs/pr/1").to(device)
+#pipe = DiffusionPipeline.from_pretrained("cocktailpeanut/xulf-s", torch_dtype=dtype).to(device)
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 2048
 
+
+def init():
+    global pipe
+
+    dtype = torch.bfloat16
+
+    # schnell is the distilled turbo model. For the CFG distilled model, use:
+    # bfl_repo = "black-forest-labs/FLUX.1-dev"
+    # revision = "refs/pr/3"
+    #
+    # The undistilled model that uses CFG ("pro") which can use negative prompts
+    # was not released.
+    bfl_repo = "cocktailpeanut/xulf-s"
+
+    scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(bfl_repo, subfolder="scheduler")
+    text_encoder = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14", torch_dtype=dtype)
+    tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14", torch_dtype=dtype)
+    text_encoder_2 = T5EncoderModel.from_pretrained(bfl_repo, subfolder="text_encoder_2", torch_dtype=dtype, revision=revision)
+    tokenizer_2 = T5TokenizerFast.from_pretrained(bfl_repo, subfolder="tokenizer_2", torch_dtype=dtype, revision=revision)
+    vae = AutoencoderKL.from_pretrained(bfl_repo, subfolder="vae", torch_dtype=dtype, revision=revision)
+    transformer = FluxTransformer2DModel.from_pretrained(bfl_repo, subfolder="transformer", torch_dtype=dtype, revision=revision)
+
+    # Experimental: Try this to load in 4-bit for <16GB cards.
+    #
+    # from optimum.quanto import qint4
+    # quantize(transformer, weights=qint4, exclude=["proj_out", "x_embedder", "norm_out", "context_embedder"])
+    # freeze(transformer)
+    quantize(transformer, weights=qfloat8)
+    freeze(transformer)
+
+    quantize(text_encoder_2, weights=qfloat8)
+    freeze(text_encoder_2)
+
+    pipe = FluxPipeline(
+        scheduler=scheduler,
+        text_encoder=text_encoder,
+        tokenizer=tokenizer,
+        text_encoder_2=None,
+        tokenizer_2=tokenizer_2,
+        vae=vae,
+        transformer=None,
+    )
+    pipe.text_encoder_2 = text_encoder_2
+    pipe.transformer = transformer
+    pipe.enable_model_cpu_offload()
+
+#    generator = torch.Generator().manual_seed(12345)
+#    image = pipe(
+#        prompt='nekomusume cat girl, digital painting',
+#        width=1024,
+#        height=1024,
+#        num_inference_steps=4,
+#        generator=generator,
+#        guidance_scale=3.5,
+#    ).images[0]
+
 #@spaces.GPU()
 def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=4, progress=gr.Progress(track_tqdm=True)):
+    global pipe
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
     generator = torch.Generator().manual_seed(seed)

requirements.txt

@@ -1,9 +1,11 @@
 accelerate
 #git+https://github.com/huggingface/diffusers.git@flux-pipeline
-git+https://github.com/peanutcocktail/diffusers.git
+#git+https://github.com/peanutcocktail/diffusers.git
+git+https://github.com/huggingface/diffusers.git
 invisible_watermark
 #torch
 transformers==4.42.4
 protobuf
 #xformers
 sentencepiece
+optimum-quanto