Update app.py

app.py

@@ -30,12 +30,37 @@ class timer:
 if not path.exists(cache_path):
     os.makedirs(cache_path, exist_ok=True)
 from huggingface_hub import hf_hub_download
+import torch
+
+from diffusers import FluxPipeline, FluxTransformer2DModel, GGUFQuantizationConfig
+
+ckpt_path = (
+    "https://huggingface.co/city96/FLUX.1-dev-gguf/blob/main/flux1-schnell-Q3_K_S.gguf"
+)
+transformer = FluxTransformer2DModel.from_single_file(
+    ckpt_path,
+    quantization_config=GGUFQuantizationConfig(compute_dtype=torch.bfloat16),
+    torch_dtype=torch.bfloat16,
+)
+pipe = FluxPipeline.from_pretrained(
+    "black-forest-labs/FLUX.1-dev",
+    transformer=transformer,
+    torch_dtype=torch.bfloat16,
+)
 
-hf_hub_download(repo_id="city96/FLUX.1-schnell-gguf", filename="flux1-schnell-Q3_K_S.gguf")
 pipe = FluxPipeline.from_pretrained("flux1-schnell-Q3_K_S.gguf")
 pipe.load_lora_weights(hf_hub_download("ByteDance/Hyper-SD", "Hyper-FLUX.1-dev-8steps-lora.safetensors"))
 pipe.fuse_lora(lora_scale=0.125)
-pipe.to(device="cuda", dtype=torch.bfloat16)
+
+pipe.enable_model_cpu_offload()
+prompt = "A cat holding a sign that says hello world"
+image = pipe(prompt, generator=torch.manual_seed(0)).images[0]
+image.save("flux-gguf.png")
+
+
+#hf_hub_download(repo_id="city96/FLUX.1-schnell-gguf", filename="flux1-schnell-Q3_K_S.gguf")
+
+pipe.to(device="cpu", dtype=torch.bfloat16)
 
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown(