Update app.py

app.py

@@ -3,8 +3,7 @@ from diffusers import StableDiffusion3Pipeline
 from huggingface_hub import login
 import os
 import gradio as gr
-from diffusers import BitsAndBytesConfig
-from diffusers import SD3Transformer2DModel
+from transformers import pipeline as transformers_pipeline
 
 # Retrieve the token from the environment variable
 token = os.getenv("HF_TOKEN")  # Hugging Face token from the secret
@@ -13,30 +12,14 @@ if token:
 else:
     raise ValueError("Hugging Face token not found. Please set it as a repository secret in the Space settings.")
 
-# Define quantization configuration (4-bit quantization)
-quant_config = BitsAndBytesConfig(
-    load_in_4bit=True,  # Enable 4-bit quantization
-    bnb_4bit_quant_type="nf4",  # Choose the quantization type (nf4 is often used for high-quality quantization)
-    bnb_4bit_compute_dtype=torch.bfloat16  # Use bfloat16 for computation (works well with CPUs)
-)
-
-# Load the Stable Diffusion 3.5 model with quantization
+# Load the Stable Diffusion 3.5 model with quantization enabled for CPU
 model_id = "stabilityai/stable-diffusion-3.5-large"
-model = SD3Transformer2DModel.from_pretrained(
-    model_id,
-    subfolder="transformer",
-    quantization_config=quant_config,
-    torch_dtype=torch.bfloat16  # Ensure the model uses bfloat16 dtype for computation
-)
+pipe = StableDiffusion3Pipeline.from_pretrained(model_id, torch_dtype=torch.float16)  # Use float16 for less memory usage
 
-# Load the pipeline with the quantized model
-pipe = StableDiffusion3Pipeline.from_pretrained(
-    model_id,
-    transformer=model,
-    torch_dtype=torch.bfloat16  # Ensuring the pipeline uses bfloat16
-)
+# Perform quantization on the model to reduce the memory footprint
+pipe.unet = torch.quantization.quantize_dynamic(pipe.unet, {torch.nn.Linear}, dtype=torch.qint8)
 
-pipe.to("cpu")  # Ensuring it runs on CPU
+pipe.to("cpu")  # Ensure it runs on CPU
 
 # Define the path to the LoRA model
 lora_model_path = "./lora_model.pth"  # Assuming the file is saved locally