image_generator / app.py
Nawinkumar15's picture
Update app.py
d38413b verified
import gradio as gr
import torch
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
# ---- Configuration ---- #
MODEL_ID = "runwayml/stable-diffusion-v1-5"
USE_CUDA = torch.cuda.is_available()
DEVICE = "cuda" if USE_CUDA else "cpu"
DTYPE = torch.float16 if USE_CUDA else torch.float32
# ---- Load Pipeline ---- #
pipe = StableDiffusionPipeline.from_pretrained(
MODEL_ID,
torch_dtype=DTYPE
).to(DEVICE)
pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
pipe.enable_attention_slicing()
# ---- Image Generation Function ---- #
def generate_image(prompt):
enhanced_prompt = f"{prompt}, ultra realistic, high detail, 8k resolution, DSLR photography, natural lighting"
with torch.inference_mode():
result = pipe(enhanced_prompt, num_inference_steps=25)
return result.images[0]
# ---- Gradio UI ---- #
demo = gr.Interface(
fn=generate_image,
inputs=gr.Textbox(
label="Enter image description",
placeholder="e.g. A cozy cabin in a snowy forest"
),
outputs=gr.Image(type="pil"),
title="🖼️ Realistic Text-to-Image Generator",
description="Generate high-quality, photorealistic images using Stable Diffusion v1.5 + DPM Scheduler"
)
# ---- Launch ---- #
if __name__ == "__main__":
demo.queue().launch()