Update app.py

app.py

@@ -3,38 +3,60 @@ import numpy as np
 import random
 from diffusers import DiffusionPipeline
 import torch
+from PIL import Image
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
+base_model = 'stabilityai/stable-diffusion-2' 
+prj_path = "bayndrysf/dreambooth-project-styl"
+huggingface_user = "bayndrysf" 
+
 if torch.cuda.is_available():
     torch.cuda.max_memory_allocated(device=device)
-    pipe = DiffusionPipeline.from_pretrained("stabilityai/sdxl-turbo", torch_dtype=torch.float16, variant="fp16", use_safetensors=True)
+    pipe = DiffusionPipeline.from_pretrained(base_model)
     pipe.enable_xformers_memory_efficient_attention()
-    pipe = pipe.to(device)
+    pipe.to(device);
+    pipe.load_lora_weights(prj_path, weight_name="pytorch_lora_weights.safetensors")
 else: 
-    pipe = DiffusionPipeline.from_pretrained("stabilityai/sdxl-turbo", use_safetensors=True)
-    pipe = pipe.to(device)
+    pipe = DiffusionPipeline.from_pretrained(base_model)
+    pipe.to(device);
+    pipe.load_lora_weights(prj_path, weight_name="pytorch_lora_weights.safetensors")
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
-def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps):
 
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-        
-    generator = torch.Generator().manual_seed(seed)
-    
-    image = pipe(
-        prompt = prompt, 
-        negative_prompt = negative_prompt,
-        guidance_scale = guidance_scale, 
-        num_inference_steps = num_inference_steps, 
-        width = width, 
-        height = height,
-        generator = generator
-    ).images[0] 
-    
+
+def image_grid(imgs, rows, cols, resize=256):
+    assert len(imgs) == rows * cols
+
+    if resize is not None:
+        imgs = [img.resize((resize, resize)) for img in imgs]
+
+    w, h = imgs[0].size
+    grid_w, grid_h = cols * w, rows * h
+    grid = Image.new("RGB", size=(grid_w, grid_h))
+
+    for i, img in enumerate(imgs):
+        x = i % cols * w
+        y = i // cols * h
+        grid.paste(img, box=(x, y))
+
+    return grid
+
+
+MAX_SEED = np.iinfo(np.int32).max
+MAX_IMAGE_SIZE = 1024
+
+
+def generate_image(prompt):
+    image = pipe(prompt=prompt, num_inference_steps=25, num_images_per_prompt = 1)
+    return image_grid(image.images, 1, 1, 512)
+
+
+
+def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps):
+    image = generate_image(prompt)
     return image
 
 examples = [