Update README.md

README.md

@@ -1,3 +1,65 @@
----
-license: apache-2.0
----
+---
+license: apache-2.0
+---
+'''python3
+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
+import torch
+from huggingface_hub import hf_hub_download
+from torchao.quantization.quant_api import (
+    quantize_,
+    int8_weight_only
+)
+dtype = torch.bfloat16
+flux_repo = "black-forest-labs/FLUX.1-schnell"
+revision = "refs/pr/1"
+
+tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14", torch_dtype=dtype)
+tokenizer_2 = T5TokenizerFast.from_pretrained(flux_repo, subfolder="tokenizer_2", torch_dtype=dtype, revision=revision)
+scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(flux_repo, subfolder="scheduler", revision=revision)
+transformer = FluxTransformer2DModel.from_pretrained(flux_repo, subfolder="transformer", torch_dtype=dtype, revision=revision)
+lora_file_path = hf_hub_download(repo_id = "Jonathan-Zhou/Flux-GameLabel-Lora", filename = "lora.safetensors")
+text_encoder = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14", torch_dtype=dtype)
+text_encoder_2 = T5EncoderModel.from_pretrained(flux_repo, subfolder="text_encoder_2", torch_dtype=dtype, revision=revision)
+vae = AutoencoderKL.from_pretrained(flux_repo, subfolder="vae", torch_dtype=dtype, revision=revision)
+
+
+pipe = FluxPipeline(
+    scheduler=scheduler,
+    text_encoder=text_encoder,
+    tokenizer=tokenizer,
+    text_encoder_2=text_encoder_2,
+    tokenizer_2=tokenizer_2,
+    vae=vae,
+    transformer=transformer,
+)
+
+# If you want to compare the lora with the bsae model, you can comment out these two lines
+pipe.load_lora_weights(lora_file_path, adapter_name="lora1")
+pipe.fuse_lora()
+
+# Quantization needed if run on a GPU with 24 GB VRAM
+quantize_(transformer, int8_weight_only()) 
+quantize_(text_encoder, int8_weight_only())
+quantize_(text_encoder_2, int8_weight_only())
+quantize_(vae, int8_weight_only())
+
+
+pipe.to("cuda")
+torch.cuda.empty_cache()
+generator = torch.Generator().manual_seed(12345)
+output = pipe(
+            prompt="a man showing off his cool new t shirt at the beach, a shark is jumping out of the water in the background", 
+            width=1024,
+            height=1024,
+            num_inference_steps=6, 
+            num_images_per_prompt = 1,
+            generator=generator,
+            guidance_scale=3.5,
+        )
+image = output.images[0]
+image.show()
+
+'''