updates 2 (#2)

- updates 2 (15c309de6c0902cfc09fca32e94107b50ade0e0e)


Co-authored-by: Aditi Baheti <Aditibaheti@users.noreply.huggingface.co>

app.py

@@ -6,17 +6,15 @@ import torch
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Replace 'path/to/your/safetensors' with the actual path to your fine-tuned model's safetensors file
-model_path = "./pytorch_lora_weights.safetensors"
+# Path to your model repository and safetensors weights
+base_model_repo = "stabilityai/stable-diffusion-3-medium-diffusers"
+lora_weights_path = "./pytorch_lora_weights.safetensors"
 
-if torch.cuda.is_available():
-    torch.cuda.max_memory_allocated(device=device)
-    pipe = DiffusionPipeline.from_pretrained(model_path, torch_dtype=torch.float16, variant="fp16", use_safetensors=True)
-    pipe.enable_xformers_memory_efficient_attention()
-    pipe = pipe.to(device)
-else: 
-    pipe = DiffusionPipeline.from_pretrained(model_path, use_safetensors=True)
-    pipe = pipe.to(device)
+# Load the base model
+pipeline = DiffusionPipeline.from_pretrained(base_model_repo, torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32)
+pipeline.load_lora_weights(lora_weights_path)
+pipeline.enable_sequential_cpu_offload()
+pipeline = pipeline.to(device)
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
@@ -25,9 +23,9 @@ def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
         
-    generator = torch.Generator().manual_seed(seed)
+    generator = torch.Generator(device=device).manual_seed(seed)
     
-    image = pipe(
+    image = pipeline(
         prompt=prompt, 
         negative_prompt=negative_prompt,
         guidance_scale=guidance_scale,