Update app.py

app.py

@@ -17,39 +17,59 @@ model_file = "weights-inkpen.ckpt"
 models = preload_models_from_standard_weights(model_file, Device)
 
 
-@spaces.GPU(duration = 242)
-def generate_image(prompt, strength, seed):
-    # Your generate function adapted to accept parameters
-    output_image = generate(
-        prompt=prompt,
-        uncond_prompt="",
-        input_image=None,
-        strength=strength,
-        do_cfg=True,
-        cfg_scale=8,
-        sampler_name="ddpm",
-        n_inference_steps=50,
-        seed=seed,
-        models=models,
-        device=Device,
-        idle_device="cpu",
-        tokenizer=tokenizer,
-    )
+@spaces.GPU(duration=242)
+def generate_image(mode, prompt, strength, seed, n_inference_steps, input_image=None):
+    if mode == "Text-to-Image":
+        # If Text-to-Image, ignore the input image
+        output_image = generate(
+            prompt=prompt,
+            uncond_prompt="",  ##negative prompt
+            input_image=None,
+            strength=strength,  ##strength for skippin steps
+            do_cfg=True,  ##cassifier free guidance
+            cfg_scale=8, 
+            sampler_name="ddpm",
+            n_inference_steps=n_inference_steps,
+            seed=seed,
+            models=models,
+            device=Device,
+            idle_device="cpu",
+            tokenizer=tokenizer,
+        )
+    elif mode == "Image-to-Image":
+        # If Image-to-Image, use the uploaded image
+        output_image = generate(
+            prompt=prompt,
+            uncond_prompt="",
+            input_image=input_image,
+            strength=strength,
+            do_cfg=True,
+            cfg_scale=8,
+            sampler_name="ddpm",
+            n_inference_steps=n_inference_steps,
+            seed=seed,
+            models=models,
+            device=Device,
+            idle_device="cpu",
+            tokenizer=tokenizer,
+        )
     return Image.fromarray(output_image)
 
-
-# Gradio interface with additional inputs for strength and seed
+# Gradio interface with additional inputs for mode selection and inference steps
 iface = gr.Interface(
     fn=generate_image,
     inputs=[
+        gr.Radio(choices=["Text-to-Image", "Image-to-Image"], label="Mode"),  # Toggle between modes
         gr.Textbox(label="Prompt"),  # Input text prompt
-        gr.Slider(0, 1, step=0.01, label="Strength (For Image-2-Image): Strength = 1 (More deviation), Strength = 0 (Less deviation)"),  # Slider for strength
+        gr.Slider(0, 1, step=0.01, label="Strength (For Image-2-Image)"),  # Slider for strength
         gr.Number(label="Seed"),  # Seed for reproducibility
+        gr.Slider(1, 1000, step=1, label="Number of Inference Steps"),  # Slider for inference steps
+        gr.Image(label="Input Image (For Image-to-Image)", optional=True),  # Optional input image
     ],
     outputs=gr.Image(label="Generated Image"),  # Output generated image
     title="Stable Diffusion Image Generator",
-    description="Generate images from text prompts using Stable Diffusion."
+    description="Generate images from text prompts or use image-to-image generation."
 )
 
 # Launch the Gradio interface with public sharing enabled
-iface.launch(share=True, debug=True)
+iface.launch(share=True, debug=True)