Spaces:
Running
on
Zero
Running
on
Zero
Avijit Ghosh
commited on
Commit
·
680331e
1
Parent(s):
b9bfe79
added SD3
Browse files
app.py
CHANGED
|
@@ -1,6 +1,6 @@
|
|
| 1 |
import gradio as gr
|
| 2 |
import torch
|
| 3 |
-
from diffusers import DiffusionPipeline, StableDiffusionPipeline, StableDiffusionXLPipeline, EulerDiscreteScheduler, UNet2DConditionModel
|
| 4 |
from transformers import BlipProcessor, BlipForConditionalGeneration
|
| 5 |
from pathlib import Path
|
| 6 |
from safetensors.torch import load_file
|
|
@@ -44,6 +44,11 @@ def load_model(model_name):
|
|
| 44 |
use_safetensors=True,
|
| 45 |
variant="fp16"
|
| 46 |
).to("cuda")
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 47 |
else:
|
| 48 |
raise ValueError("Unknown model name")
|
| 49 |
return pipeline
|
|
@@ -54,6 +59,8 @@ pipeline_text2image = load_model(default_model)
|
|
| 54 |
|
| 55 |
@spaces.GPU
|
| 56 |
def getimgen(prompt, model_name):
|
|
|
|
|
|
|
| 57 |
if model_name == "stabilityai/sdxl-turbo":
|
| 58 |
return pipeline_text2image(prompt=prompt, guidance_scale=0.0, num_inference_steps=2).images[0]
|
| 59 |
elif model_name == "runwayml/stable-diffusion-v1-5":
|
|
@@ -63,6 +70,8 @@ def getimgen(prompt, model_name):
|
|
| 63 |
elif model_name == "segmind/SSD-1B":
|
| 64 |
neg_prompt = "ugly, blurry, poor quality"
|
| 65 |
return pipeline_text2image(prompt=prompt, negative_prompt=neg_prompt).images[0]
|
|
|
|
|
|
|
| 66 |
|
| 67 |
blip_processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-large")
|
| 68 |
blip_model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-large", torch_dtype=torch.float16).to("cuda")
|
|
|
|
| 1 |
import gradio as gr
|
| 2 |
import torch
|
| 3 |
+
from diffusers import DiffusionPipeline, StableDiffusionPipeline, StableDiffusionXLPipeline, EulerDiscreteScheduler, UNet2DConditionModel, StableDiffusion3Pipeline
|
| 4 |
from transformers import BlipProcessor, BlipForConditionalGeneration
|
| 5 |
from pathlib import Path
|
| 6 |
from safetensors.torch import load_file
|
|
|
|
| 44 |
use_safetensors=True,
|
| 45 |
variant="fp16"
|
| 46 |
).to("cuda")
|
| 47 |
+
elif model_name == "stabilityai/stable-diffusion-3-medium-diffusers":
|
| 48 |
+
pipeline = StableDiffusion3Pipeline.from_pretrained(
|
| 49 |
+
model_name,
|
| 50 |
+
torch_dtype=torch.float16
|
| 51 |
+
).to("cuda")
|
| 52 |
else:
|
| 53 |
raise ValueError("Unknown model name")
|
| 54 |
return pipeline
|
|
|
|
| 59 |
|
| 60 |
@spaces.GPU
|
| 61 |
def getimgen(prompt, model_name):
|
| 62 |
+
global pipeline_text2image
|
| 63 |
+
pipeline_text2image = load_model(model_name)
|
| 64 |
if model_name == "stabilityai/sdxl-turbo":
|
| 65 |
return pipeline_text2image(prompt=prompt, guidance_scale=0.0, num_inference_steps=2).images[0]
|
| 66 |
elif model_name == "runwayml/stable-diffusion-v1-5":
|
|
|
|
| 70 |
elif model_name == "segmind/SSD-1B":
|
| 71 |
neg_prompt = "ugly, blurry, poor quality"
|
| 72 |
return pipeline_text2image(prompt=prompt, negative_prompt=neg_prompt).images[0]
|
| 73 |
+
elif model_name == "stabilityai/stable-diffusion-3-medium-diffusers":
|
| 74 |
+
return pipeline_text2image(prompt=prompt, negative_prompt="", num_inference_steps=28, guidance_scale=7.0).images[0]
|
| 75 |
|
| 76 |
blip_processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-large")
|
| 77 |
blip_model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-large", torch_dtype=torch.float16).to("cuda")
|