Spaces:
Runtime error
Runtime error
File size: 1,958 Bytes
a9c09c1 df9c0f8 a9c09c1 86ba726 a9c09c1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 |
from diffusers import DDPMPipeline
import torch
import numpy as np
import gradio as gr
from torchvision.utils import make_grid
import torchvision.transforms as transforms
from PIL import Image
from transformers import logging
from transformers.utils import move_cache
# Optional: handle potential cache migration for Transformers
try:
move_cache()
except Exception as e:
logging.error(f"Error migrating cache: {str(e)}")
# Check for 'accelerate' library and suggest installation if not found
try:
from accelerate import Accelerator
except ImportError:
logging.warning("Accelerate library not found. It's recommended to install it for efficient model loading.")
# Load the pre-trained pipeline from Hugging Face
MODEL_ID = "ahmetyaylalioglu/textile_diffusion_ddpm"
pipeline = DDPMPipeline.from_pretrained(MODEL_ID)
pipeline.to("cuda")
pipeline.unet.eval()
def generate_images(seed, num_images):
# Set the seed
seed = int(seed) # Ensure the seed is an integer
torch.manual_seed(seed)
# Generate images
num_images = int(num_images) # Ensure num_images is an integer
generated_images = pipeline(batch_size=num_images).images
# Convert the list of generated images to a grid
image_grid = make_grid([transforms.ToTensor()(img) for img in generated_images], nrow=4)
# Convert the grid to a PIL image
image_grid_pil = transforms.ToPILImage()(image_grid)
return image_grid_pil
# Set up the Gradio interface
interface = gr.Interface(
fn=generate_images,
inputs=[
gr.Textbox(label="Random Seed", placeholder="Enter a seed number", value=str(np.random.randint(0, 1000))),
gr.Textbox(label="Number of Images", placeholder="Enter number of images to generate", value="8")
],
outputs="image",
title="Textile Diffusion DDPM",
description="Generate textile images using a trained DDPM model from Hugging Face."
)
# Launch the interface
interface.launch()
|