Update app.py

app.py

@@ -71,12 +71,19 @@ def get_images(ws, prompt, server_address):
     return output_images
 
 # 🌟 Chapter 7: Generate Image Function 🌟
-def generate_image(text_prompt, seed, server):
-    prompt_text = """
-    {
-        "3": {
+def generate_image(text_prompt, seed, quality, server):
+    steps_map = {
+        "Low": 8,
+        "Medium": 16,
+        "High": 30
+    }
+    steps = steps_map[quality]
+
+    prompt_text = f"""
+    {{
+        "3": {{
             "class_type": "KSampler",
-            "inputs": {
+            "inputs": {{
                 "cfg": 8,
                 "denoise": 1,
                 "latent_image": [
@@ -97,47 +104,47 @@ def generate_image(text_prompt, seed, server):
                 ],
                 "sampler_name": "euler",
                 "scheduler": "normal",
-                "seed": 8566257,
-                "steps": 8
-            }
-        },
-        "4": {
+                "seed": {seed},
+                "steps": {steps}
+            }}
+        }},
+        "4": {{
             "class_type": "CheckpointLoaderSimple",
-            "inputs": {
+            "inputs": {{
                 "ckpt_name": "v1-5-pruned-emaonly.ckpt"
-            }
-        },
-        "5": {
+            }}
+        }},
+        "5": {{
             "class_type": "EmptyLatentImage",
-            "inputs": {
+            "inputs": {{
                 "batch_size": 1,
                 "height": 512,
                 "width": 768
-            }
-        },
-        "6": {
+            }}
+        }},
+        "6": {{
             "class_type": "CLIPTextEncode",
-            "inputs": {
+            "inputs": {{
                 "clip": [
                     "4",
                     1
                 ],
-                "text": "masterpiece best quality girl"
-            }
-        },
-        "7": {
+                "text": "{text_prompt}"
+            }}
+        }},
+        "7": {{
             "class_type": "CLIPTextEncode",
-            "inputs": {
+            "inputs": {{
                 "clip": [
                     "4",
                     1
                 ],
                 "text": "bad hands"
-            }
-        },
-        "8": {
+            }}
+        }},
+        "8": {{
             "class_type": "VAEDecode",
-            "inputs": {
+            "inputs": {{
                 "samples": [
                     "3",
                     0
@@ -146,23 +153,21 @@ def generate_image(text_prompt, seed, server):
                     "4",
                     2
                 ]
-            }
-        },
-        "save_image_websocket_node": {
+            }}
+        }},
+        "save_image_websocket_node": {{
             "class_type": "SaveImageWebsocket",
-            "inputs": {
+            "inputs": {{
                 "images": [
                     "8",
                     0
                 ]
-            }
-        }
-    }
+            }}
+        }}
+    }}
     """
 
     prompt = json.loads(prompt_text)
-    prompt["6"]["inputs"]["text"] = text_prompt
-    prompt["3"]["inputs"]["seed"] = seed
     server_address = "3.14.144.23:8188" if server == "AWS Server" else "192.168.50.136:8188"
     ws = websocket.WebSocket()
     ws.connect(f"ws://{server_address}/ws?clientId={client_id}")
@@ -191,6 +196,7 @@ with gr.Blocks() as demo:
         with gr.Column():
             text_prompt = gr.Textbox(label="Text Prompt", value="masterpiece best quality man")
             seed = gr.Number(label="Seed", value=5)
+            quality = gr.Radio(label="Quality", choices=["Low", "Medium", "High"], value="Low")
             server = gr.Radio(label="Server", choices=["AWS Server", "Home Server"], value="AWS Server")
             generate_button = gr.Button("Generate Image")
             cancel_button = gr.Button("Cancel Request")
@@ -198,7 +204,7 @@ with gr.Blocks() as demo:
         with gr.Column():
             output_image = gr.Image(label="Generated Image")
     
-    generate_button.click(fn=generate_image, inputs=[text_prompt, seed, server], outputs=output_image)
+    generate_button.click(fn=generate_image, inputs=[text_prompt, seed, quality, server], outputs=output_image)
     cancel_button.click(fn=cancel_request, inputs=[], outputs=[])
 
 demo.launch()