Update README.md

README.md

@@ -70,43 +70,51 @@ You can run the smashed model with these steps:
 2. Load & run the model.
     ```python 
     
-      import torch
-      
-      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
-      
-      dtype = torch.bfloat16
-      
-      bfl_repo = "black-forest-labs/FLUX.1-schnell"
-      revision = "refs/pr/1"
-      local_path = "FLUX.1-schnell-8bit"
-      
-      scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(bfl_repo, subfolder="scheduler", revision=revision)
-      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 = torch.load(local_path + '/text_encoder_2.pt')
-      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 = torch.load(local_path + '/transformer.pt')
-      
-      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)
+    import torch
+    
+    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
+    
+    dtype = torch.bfloat16
+    
+    bfl_repo = "black-forest-labs/FLUX.1-schnell"
+    revision = "refs/pr/1"
+    local_path = "FLUX.1-schnell-8bit"
+    
+    scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(bfl_repo, subfolder="scheduler", revision=revision)
+    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 = torch.load(local_path + '/text_encoder_2.pt')
+    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 = torch.load(local_path + '/transformer.pt')
+    
+    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,
+    guidance_scale=0.0,
+    num_inference_steps=4,
+    max_sequence_length=256,
+    generator=torch.Generator("cpu").manual_seed(0)
+      ).images[0]
+    image.save("flux-schnell.png")
     ```
 
 ## Configurations