Update app.py

app.py

@@ -1,17 +1,21 @@
 import gradio as gr
 import torch
-from diffusers import StableDiffusionXLPipeline, AutoencoderKL, KDPM2AncestralDiscreteScheduler
+from diffusers import StableDiffusionXLPipeline, AutoencoderKL, KDPM2AncestralDiscreteScheduler, UNet2DConditionModel
 from huggingface_hub import hf_hub_download
 import spaces
 from PIL import Image
 import requests
 from translatepy import Translator
+import numpy as np
+import random
+
 
 translator = Translator()
 
 # Constants
 model = "Corcelio/mobius"
 vae_model = "madebyollin/sdxl-vae-fp16-fix"
+MAX_SEED = np.iinfo(np.int32).max
 
 CSS = """
 .gradio-container {
@@ -37,7 +41,8 @@ vae = AutoencoderKL.from_pretrained(
 
 # Ensure model and scheduler are initialized in GPU-enabled function
 if torch.cuda.is_available():
-    pipe = StableDiffusionXLPipeline.from_pretrained(model, vae=vae, torch_dtype=torch.float16).to("cuda")
+    unet = UNet2DConditionModel.from_pretrained(model, subfolder="unet").to("cuda", torch.float16)
+    pipe = StableDiffusionXLPipeline.from_pretrained(model, vae=vae, unet=unet, torch_dtype=torch.float16).to("cuda")
 
 pipe.scheduler = KDPM2AncestralDiscreteScheduler.from_config(pipe.scheduler.config)
 
@@ -49,13 +54,21 @@ def generate_image(
     negative="low quality",
     width=1024,
     height=1024,
+    seed=-1,
+    nums=1,
     scale=1.5,
     steps=30,
     clip=3):
     
+    if seed == -1:
+        seed = random.randint(0, MAX_SEED)
+        
+    generator = torch.Generator().manual_seed(seed)
+    
     prompt = str(translator.translate(prompt, 'English'))
 
     print(f'prompt:{prompt}')
+
     
     image = pipe(
         prompt, 
@@ -63,10 +76,12 @@ def generate_image(
         width=width,
         height=height,
         guidance_scale=scale,
+        generator = generator,
         num_inference_steps=steps,
+        num_images_per_prompt=nums,
         clip_skip=clip,
-    )
-    return image.images[0]
+    ).images
+    return image, seed
 
 
 examples = [
@@ -84,15 +99,15 @@ examples = [
 
 with gr.Blocks(css=CSS, js=JS, theme="soft") as demo:
     gr.HTML("<h1><center>Mobius💠</center></h1>")
-    gr.HTML("<p><center><a href='https://huggingface.co/Corcelio/mobius'>mobius</a> text-to-image generation</center><br><center>Multi-Languages. Adding default prompts to enhance.</center></p>")
+    gr.HTML("<p><center><a href='https://huggingface.co/Corcelio/mobius'>mobius</a> text-to-image generation</center><br><center>Adding default prompts to enhance.</center></p>")
     with gr.Group():
         with gr.Row():
-            prompt = gr.Textbox(label='Enter Your Prompt', value="best quality, HD, aesthetic", scale=6)
+            prompt = gr.Textbox(label='Enter Your Prompt(Multi-Languages)', value="best quality, HD, aesthetic", scale=6)
             submit = gr.Button(scale=1, variant='primary')
-    img = gr.Image(label='Mobius Generated Image')
+    img = gr.Gallery(label='Mobius Generated Image',columns = 1, preview=True)
     with gr.Accordion("Advanced Options", open=False):
         with gr.Row():
-            negative = gr.Textbox(label="Negative prompt", value="low quality")
+            negative = gr.Textbox(label="Negative prompt", value="low quality, ugly, blurry, poor face, bad anatomy")
         with gr.Row():
             width = gr.Slider(
                 label="Width",
@@ -108,6 +123,23 @@ with gr.Blocks(css=CSS, js=JS, theme="soft") as demo:
                 step=8,
                 value=1024,
             )
+        with gr.Row():
+            seed = gr.Slider(
+                label="Seed (-1 Get Random)",
+                minimum=-1,
+                maximum=MAX_SEED,
+                step=1,
+                value=-1,
+                scale=2,
+            )
+            nums = gr.Slider(
+                label="Image Numbers",
+                minimum=1,
+                maximum=4,
+                step=1,
+                value=1,
+                scale=1,
+            ) 
         with gr.Row():
             scale = gr.Slider(
                 label="Guidance",
@@ -129,7 +161,7 @@ with gr.Blocks(css=CSS, js=JS, theme="soft") as demo:
                 maximum=10,
                 step=1,
                 value=3,
-            )      
+            )
     gr.Examples(
         examples=examples,
         inputs=prompt,
@@ -139,11 +171,11 @@ with gr.Blocks(css=CSS, js=JS, theme="soft") as demo:
     )
 
     prompt.submit(fn=generate_image,
-                 inputs=[prompt, negative, width, height, scale, steps, clip],
+                 inputs=[prompt, negative, width, height, seed, nums, scale, steps, clip],
                  outputs=img,
                  )
     submit.click(fn=generate_image,
-                 inputs=[prompt, negative, width, height, scale, steps, clip],
+                 inputs=[prompt, negative, width, height, seed, nums, scale, steps, clip],
                  outputs=img,
                  )