Update app.py

app.py

@@ -36,11 +36,32 @@ def genie (Model, Prompt, negative_prompt, height, width, scale, steps, seed):
         return image
 
     if Model == "SD3.5 Turbo":
-        progress=gr.Progress(track_tqdm=True)
+        torch.cuda.max_memory_allocated(device=device)
         torch.cuda.empty_cache()
-        model_repo_id = "stabilityai/stable-diffusion-3.5-large-turbo"
-        pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype = torch.float32)
-        pipe = pipe.to(device)
+        progress=gr.Progress(track_tqdm=True)
+        from diffusers import BitsAndBytesConfig, SD3Transformer2DModel
+        from diffusers import StableDiffusion3Pipeline
+
+
+        model_id = "stabilityai/stable-diffusion-3.5-large-turbo"
+
+        nf4_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_compute_dtype=torch.bfloat16)
+        model_nf4 = SD3Transformer2DModel.from_pretrained(
+            model_id,
+            subfolder="transformer",
+            quantization_config=nf4_config,
+            torch_dtype=torch.bfloat16)
+
+        t5_nf4 = T5EncoderModel.from_pretrained("diffusers/t5-nf4", torch_dtype=torch.bfloat16)
+
+        pipe = StableDiffusion3Pipeline.from_pretrained(
+            model_id, 
+            transformer=model_nf4,
+            text_encoder_3=t5_nf4,
+            torch_dtype=torch.bfloat16)
         pipe.enable_model_cpu_offload()
         pipe.vae.enable_slicing()
         pipe.vae.enable_tiling()