Upload klein.py

klein.py

@@ -0,0 +1,42 @@
+import torch
+from diffusers import Flux2KleinPipeline
+from transformers import BitsAndBytesConfig
+
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_quant_type="nf4",          # BEST quality/speed
+    bnb_4bit_compute_dtype=torch.bfloat16,  # fast on Ampere+
+    bnb_4bit_use_double_quant=True,     # lower VRAM
+)
+
+device = "cuda"
+dtype = torch.bfloat16
+
+pipe = Flux2KleinPipeline.from_pretrained("./FLUX.2-9B-bnb-4bit", torch_dtype=dtype)
+
+"""
+pipe = Flux2KleinPipeline.from_pretrained(
+    "black-forest-labs/FLUX.2-klein-4B",
+    torch_dtype=torch.bfloat16,
+    device_map="auto",                 # REQUIRED
+    quantization_config=bnb_config,    # APPLY 4-bit
+)
+"""
+pipe.to("cuda")
+#pipe.enable_model_cpu_offload()  # save some VRAM by offloading the model to CPU
+
+from PIL import Image
+
+init_image = Image.open("suji.jpg").convert("RGB")
+#prompt = "an very beautiful sexy korean kpop young woman with white bikini is smiling on the waikiki beach. hiqh quality realistic photo."# pixar 3d style"
+prompt = "beautiful young korea sexy woman in the beach holding plate with Circulus "
+image = pipe(
+    prompt=prompt,
+    image=init_image,
+    height=1024,
+    width=1024,
+    guidance_scale=1.0,
+    num_inference_steps=4,
+    generator=torch.Generator(device=device).manual_seed(0)
+).images[0]
+image.save("./output/flux_suji3.png")