Update app.py

app.py

@@ -1,7 +1,6 @@
 from diffusers import AutoPipelineForText2Image, StableDiffusionImg2ImgPipeline
 from PIL import Image
 import gradio as gr
-import numpy as np
 import random
 import torch
 import math
@@ -17,36 +16,33 @@ css = """
 }
 """
 
-def generate(prompt, samp_steps, seed, strength, progress=gr.Progress(track_tqdm=True)):
+def generate(prompt, samp_steps, seed, progress=gr.Progress(track_tqdm=True)):
     if seed < 0:
         seed = random.randint(1,999999)
-#    image = txt2img(
-#        prompt,
-#       num_inference_steps=1,
-#        guidance_scale=0.0,
-#        generator=torch.manual_seed(seed),
-#    ).images[0]
-    data=np.random.randint(low=0,high=256,size=128*128*3, dtype=np.uint8)
-    data=data.reshape(128,128,3)
-    image = Image.fromarray(data,'RGB')
+    image = txt2img(
+        prompt,
+        num_inference_steps=1,
+        guidance_scale=0.0,
+        generator=torch.manual_seed(seed),
+    ).images[0]
     upscaled_image = image.resize((1024,1024), 1)
     final_image = img2img(
         prompt,
         upscaled_image,
         num_inference_steps=math.ceil(samp_steps/strength),
         guidance_scale=5,
-        strength=strength,
+        strength=1,
         generator=torch.manual_seed(seed),
     ).images[0]
     return [final_image], seed
         
 def set_base_models():
-#    txt2img = AutoPipelineForText2Image.from_pretrained(
-#        "stabilityai/sdxl-turbo",
-#        torch_dtype = torch.float16,
-#        variant = "fp16"
-#    )
-#    txt2img.to("cuda")
+    txt2img = AutoPipelineForText2Image.from_pretrained(
+        "stabilityai/sdxl-turbo",
+        torch_dtype = torch.float16,
+        variant = "fp16"
+    )
+    txt2img.to("cuda")
     img2img = StableDiffusionImg2ImgPipeline.from_pretrained(
         "Lykon/dreamshaper-8",
         torch_dtype = torch.float16,
@@ -62,14 +58,13 @@ with gr.Blocks(css=css) as demo:
         submit_btn = gr.Button("Generate", elem_classes="btn-green")
         
         with gr.Row():
-            sampling_steps = gr.Slider(1, 10, value=5, step=1, label="Sampling steps")
-            strength = gr.Slider(0, 1, value=0.75, step=0.05, label="Refiner strength")
+            sampling_steps = gr.Slider(1, 6, value=3, step=1, label="Refiner steps")
             seed = gr.Number(label="Seed", value=-1, minimum=-1, precision=0)
             lastSeed = gr.Number(label="Last Seed", value=-1, interactive=False)
             
         gallery = gr.Gallery(show_label=False, preview=True, container=False, height=1100)
         
-    submit_btn.click(generate, [prompt, sampling_steps, seed, strength], [gallery, lastSeed], queue=True)
+    submit_btn.click(generate, [prompt, sampling_steps, seed], [gallery, lastSeed], queue=True)
     
 txt2img, img2img = set_base_models()
 demo.launch(debug=True)