Update README.md

app.py

@@ -17,10 +17,10 @@ pipe, params = FlaxStableDiffusionPipeline.from_pretrained(
     use_memory_efficient_attention=True
 )
 
-def infer(prompts, negative_prompts, width=1088, height=1088, inference_steps=30, seed=0):
+def infer(prompts, negative_prompts):
     
     num_samples = 1 #jax.device_count()
-    rng = create_key(int(seed))
+    rng = create_key(0)
     rng = jax.random.split(rng, jax.device_count())
     
     prompt_ids = pipe.prepare_inputs([prompts] * num_samples)
@@ -33,57 +33,15 @@ def infer(prompts, negative_prompts, width=1088, height=1088, inference_steps=30
     output = pipe(
         prompt_ids=prompt_ids,
         params=p_params,
-        height=height,
-        width=width,
+        height=1088,
+        width=1088,
         prng_seed=rng,
-        num_inference_steps=inference_steps,
+        num_inference_steps=50,
         neg_prompt_ids=negative_prompt_ids,
         jit=True,
     ).images
     
     output_images = pipe.numpy_to_pil(np.asarray(output.reshape((num_samples,) + output.shape[-3:])))
-    return output_images[0]
+    return output_images
 
-prompt_input = gr.inputs.Textbox(
-    label="Prompt", 
-    placeholder="a highly detailed mansion in the autumn by studio ghibli, makoto shinkai"
-)
-neg_prompt_input = gr.inputs.Textbox(
-    label="Negative Prompt", 
-    placeholder=""
-)
-
-width_slider = gr.inputs.Slider(
-    minimum=512, maximum=2048, default=1088, step=64, label="width"
-)
-
-height_slider = gr.inputs.Slider(
-    minimum=512, maximum=2048, default=1088, step=64, label="height"
-)
-
-inf_steps_input = gr.inputs.Slider(
-    minimum=1, maximum=100, default=30, step=1, label="Inference Steps"
-)
-
-
-seed_input = gr.inputs.Number(default=0, label="Seed")
-
-app = gr.Interface(
-    fn=infer,
-    inputs=[prompt_input, neg_prompt_input, width_slider, height_slider, inf_steps_input, seed_input],
-    outputs="image",
-    title="Stable Diffusion High Resolution",
-    description=(
-        "Based on stable diffusion 1.5 and fine-tuned on 576x576 up to 1088x1088 images, "
-        "Stable Diffusion High Resolution is compartible with another SD1.5 model and mergeable with other SD1.5 model, "
-        "giving other model to generate high resolution images without using upscaler."
-    ),
-    examples=[
-        ["a highly detailed mansion in the autumn by studio ghibli, makoto shinkai","", 1088, 1088, 30, 0],
-        ["best high quality landscape, in the morning light, Overlooking TOKYO beautiful city with Fujiyama, from a tall house, by greg rutkowski and thomas kinkade, Trending on artstation makoto shinkai style","", 1088, 576, 30, 0],
-        [" assassin's creed black flag, hd, 4k, dlsr ","", 960, 960, 30, 4154731],
-    ],
-    
-)
-
-app.launch()
+gr.Interface(infer, inputs=["text", "text"], outputs="gallery").launch()