from diffusers import StableDiffusionPipeline
import torch
from PIL import Image
import math
import gradio as gr

pipeline = StableDiffusionPipeline.from_pretrained("MohamedRashad/diffusion_fashion")
#pipeline.to("cuda")

def image_grid(imgs, rows, cols):
    assert len(imgs) == rows*cols

    w, h = imgs[0].size
    grid = Image.new('RGB', size=(cols*w, rows*h))
    grid_w, grid_h = grid.size

    for i, img in enumerate(imgs):
        grid.paste(img, box=(i%cols*w, i//cols*h))
    return grid


def generate_image2(prompt,num_images):
 # num_images = 3
  num_images=math.floor(num_images)
  prompt = [prompt] * num_images
  #r=num_images//2
  #c=2
  images = pipeline(prompt).images
  #print(images)
  grid = image_grid(images, rows=1, cols=num_images)

  grid.save(f"prompt.png")
  return grid


text_input = gr.inputs.Textbox(label="Enter prompt")
number_input = gr.inputs.Number(label="Enter Number of images")

demo2=gr.Interface(
    fn=generate_image2,
    inputs=[text_input, number_input],
    outputs="image",
    title="Image Generation",
    description="Enter a prompt and see a grid of generated images.",
    layout="vertical",
)

demo2.launch(inline=False)