Add SD2

app.py

@@ -1,6 +1,13 @@
 import gradio as gr
 import torch
-from diffusers import DiffusionPipeline, StableDiffusionPipeline, StableDiffusionXLPipeline, EulerDiscreteScheduler, UNet2DConditionModel, StableDiffusion3Pipeline
+from diffusers import (
+    DiffusionPipeline,
+    StableDiffusionPipeline,
+    StableDiffusionXLPipeline,
+    EulerDiscreteScheduler,
+    UNet2DConditionModel,
+    StableDiffusion3Pipeline
+)
 from transformers import BlipProcessor, BlipForConditionalGeneration
 from pathlib import Path
 from safetensors.torch import load_file
@@ -54,7 +61,14 @@ def load_model(model_name):
     elif model_name == "stabilityai/stable-diffusion-3-medium-diffusers":
         pipeline = StableDiffusion3Pipeline.from_pretrained(
             model_name, 
-            torch_dtype=torch.float16,
+            torch_dtype=torch.float16
+        ).to("cuda")
+    elif model_name == "stabilityai/stable-diffusion-2":
+        scheduler = EulerDiscreteScheduler.from_pretrained(model_name, subfolder="scheduler")
+        pipeline = StableDiffusionPipeline.from_pretrained(
+            model_name, 
+            scheduler=scheduler, 
+            torch_dtype=torch.float16
         ).to("cuda")
     else:
         raise ValueError("Unknown model name")
@@ -77,6 +91,8 @@ def getimgen(prompt, model_name):
         return pipeline_text2image(prompt=prompt, negative_prompt=neg_prompt).images[0]
     elif model_name == "stabilityai/stable-diffusion-3-medium-diffusers":
         return pipeline_text2image(prompt=prompt, negative_prompt="", num_inference_steps=28, guidance_scale=7.0).images[0]
+    elif model_name == "stabilityai/stable-diffusion-2":
+        return pipeline_text2image(prompt=prompt).images[0]
 
 blip_processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-large")
 blip_model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-large", torch_dtype=torch.float16).to("cuda")
@@ -167,7 +183,8 @@ This demo provides an insightful look into how current text-to-image models hand
         choices=[
             "stabilityai/stable-diffusion-3-medium-diffusers",
             "stabilityai/sdxl-turbo", 
-            "ByteDance/SDXL-Lightning", 
+            "ByteDance/SDXL-Lightning",
+            "stabilityai/stable-diffusion-2",
             "runwayml/stable-diffusion-v1-5", 
             "segmind/SSD-1B"
         ],