Avijit Ghosh commited on
Commit
680331e
1 Parent(s): b9bfe79
Files changed (1) hide show
  1. app.py +10 -1
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")