text2image22 / app.py
Rahul8827's picture
Update app.py
fb1a904
import torch
import gradio as gr
from diffusers import StableDiffusionPipeline
class CFG:
device = "cpu"
seed = 42
generator = torch.Generator(device).manual_seed(seed)
image_gen_steps = 35
image_gen_model_id = "stabilityai/stable-diffusion-2"
image_gen_size = (400, 400)
image_gen_guidance_scale = 9
image_gen_model = StableDiffusionPipeline.from_pretrained(
CFG.image_gen_model_id, torch_dtype=torch.float32,
revision="fp16", use_auth_token='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', guidance_scale=9
)
apply = image_gen_model.to(CFG.device)
def generate_image(prompt):
image = apply(
prompt, num_inference_steps=CFG.image_gen_steps,
generator=CFG.generator,
guidance_scale=CFG.image_gen_guidance_scale
).images[0]
image = image.resize(CFG.image_gen_size)
return image
title = "Stable Diffusion Image Generator"
description = "Enter a prompt to generate an image using the Stable Diffusion model."
iface = gr.Interface(fn=generate_image, inputs="text", outputs="image", title=title, description=description)
iface.launch()