|
|
|
|
|
|
|
|
import gradio as gr |
|
|
from diffusers import StableDiffusionPipeline |
|
|
import torch |
|
|
from PIL import Image |
|
|
import traceback |
|
|
|
|
|
|
|
|
MODEL_ID = "runwayml/stable-diffusion-v1-5" |
|
|
|
|
|
|
|
|
try: |
|
|
device = "cuda" if torch.cuda.is_available() else "cpu" |
|
|
dtype = torch.float16 if torch.cuda.is_available() else torch.float32 |
|
|
|
|
|
pipe = StableDiffusionPipeline.from_pretrained( |
|
|
MODEL_ID, |
|
|
torch_dtype=dtype, |
|
|
safety_checker=None |
|
|
).to(device) |
|
|
|
|
|
except Exception as e: |
|
|
print("❌ Model Loading Error:") |
|
|
traceback.print_exc() |
|
|
pipe = None |
|
|
|
|
|
|
|
|
def generate_image(prompt): |
|
|
if pipe is None: |
|
|
return "⚠️ Model not loaded properly. Please check logs." |
|
|
try: |
|
|
result = pipe(prompt) |
|
|
image = result.images[0] |
|
|
return image |
|
|
except Exception as e: |
|
|
print("❌ Generation Error:", e) |
|
|
return "⚠️ Error generating image." |
|
|
|
|
|
|
|
|
title = "🪷 Shrividya Text-to-Image AI" |
|
|
description = "Generate stunning images from your text prompts using Stable Diffusion." |
|
|
|
|
|
iface = gr.Interface( |
|
|
fn=generate_image, |
|
|
inputs=gr.Textbox(label="Enter your prompt", placeholder="e.g. Divine temple on riverbank at sunset"), |
|
|
outputs=gr.Image(label="Generated Image"), |
|
|
title=title, |
|
|
description=description, |
|
|
allow_flagging="never" |
|
|
) |
|
|
|
|
|
if __name__ == "__main__": |
|
|
iface.launch() |
|
|
|