Update app.py

app.py

@@ -1,16 +1,132 @@
+# Ensure necessary packages are installed
+!pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118  # Install specific version of torchvision+ 
+!pip install git+https://github.com/huggingface/diffusers.git  # Install latest diffusers
+!pip install torch gradio basicsr realesrgan Pillow diffusers transformers accelerate safetensors  # Install dependencies
+
+# Import necessary libraries
 import torch
-from diffusers import StableDiffusionPipeline
 import gradio as gr
+import webbrowser
+from diffusers import StableDiffusionPipeline, StableDiffusionImg2ImgPipeline
+from PIL import Image
+
+# Check if GPU is available
+device = "cuda" if torch.cuda.is_available() else "cpu"
+if device == "cpu":
+    print("⚠️ Warning: Running on CPU, performance may be slow.")
+
+# Load Text-to-Image model
+print("🔄 Loading Stable Diffusion txt2img model...")
+pipe_txt2img = StableDiffusionPipeline.from_pretrained(
+    "runwayml/stable-diffusion-v1-5",
+    torch_dtype=torch.float16 if device == "cuda" else torch.float32
+).to(device)
+print("✅ Text-to-Image model loaded!")
+
+# Load Image-to-Image model
+print("🔄 Loading Stable Diffusion img2img model...")
+pipe_img2img = StableDiffusionImg2ImgPipeline.from_pretrained(
+    "runwayml/stable-diffusion-v1-5",
+    torch_dtype=torch.float16 if device == "cuda" else torch.float32
+).to(device)
+print("✅ Image-to-Image model loaded!")
+
+# Function to generate images from text
+def generate_txt2img(prompt, steps=50, guidance=7.5, width=512, height=512, seed=-1, save_format="png"):
+    generator = torch.manual_seed(seed) if seed != -1 else None
+    image = pipe_txt2img(
+        prompt, num_inference_steps=steps, guidance_scale=guidance, width=width, height=height,
+        generator=generator
+    ).images[0]
+    
+    output_path = f"generated_image.{save_format}"  # Save image in the requested format
+    image.save(output_path, format=save_format.upper())  # Save in the selected format
+    print(f"Image saved to {output_path}")
+    return output_path
+
+# Function to transform images using img2img
+def generate_img2img(prompt, image, strength=0.5, steps=50, guidance=7.5, width=512, height=512, seed=-1, save_format="png"):
+    generator = torch.manual_seed(seed) if seed != -1 else None
+    image = pipe_img2img(
+        prompt, image=image, strength=strength, num_inference_steps=steps, guidance_scale=guidance,
+        width=width, height=height, generator=generator
+    ).images[0]
+    
+    output_path = f"modified_image.{save_format}"  # Save image in the requested format
+    image.save(output_path, format=save_format.upper())  # Save in the selected format
+    print(f"Image saved to {output_path}")
+    return output_path
+
+# Define Gradio UI
+def create_ui():
+    with gr.Blocks(title="DiffuGen: AI Image Generation") as demo:
+        gr.Markdown("# 🌟 DiffuGen - AI Image Generator")
+
+        # Text-to-Image Tab
+        with gr.Tab("📷 Text to Image"):
+            with gr.Row():
+                prompt = gr.Textbox(label="Enter a text prompt")
+
+            with gr.Row():
+                steps = gr.Slider(10, 100, value=50, step=10, label="Steps")
+                guidance = gr.Slider(1, 15, value=7.5, label="Guidance Scale")
+
+            with gr.Row():
+                width = gr.Slider(256, 1024, value=512, step=64, label="Width")
+                height = gr.Slider(256, 1024, value=512, step=64, label="Height")
+                seed = gr.Number(value=-1, label="Seed (-1 for random)")
+
+            with gr.Row():
+                save_format = gr.Dropdown(
+                    choices=["png", "jpg"], value="png", label="Select Image Format"
+                )
+
+            generate_btn = gr.Button("🚀 Generate Image")
+            output_image = gr.Image(label="Generated Image", type="pil")
+
+            generate_btn.click(
+                generate_txt2img, 
+                inputs=[prompt, steps, guidance, width, height, seed, save_format], 
+                outputs=output_image
+            )
+
+        # Image-to-Image Tab
+        with gr.Tab("🖼️ Image to Image"):
+            with gr.Row():
+                prompt_img2img = gr.Textbox(label="Enter a prompt")
+
+            with gr.Row():
+                input_img = gr.Image(label="Upload Image", type="pil")
+
+            with gr.Row():
+                strength = gr.Slider(0.1, 1.0, value=0.5, label="Denoising Strength")
+                steps_img2img = gr.Slider(10, 100, value=50, label="Steps")
+                guidance_img2img = gr.Slider(1, 15, value=7.5, label="Guidance Scale")
+
+            with gr.Row():
+                width_img2img = gr.Slider(256, 1024, value=512, step=64, label="Width")
+                height_img2img = gr.Slider(256, 1024, value=512, step=64, label="Height")
+                seed_img2img = gr.Number(value=-1, label="Seed (-1 for random)")
+
+            with gr.Row():
+                save_format_img2img = gr.Dropdown(
+                    choices=["png", "jpg"], value="png", label="Select Image Format"
+                )
+
+            generate_img_btn = gr.Button("🔄 Transform Image")
+            output_img2img = gr.Image(label="Modified Image", type="pil")
+
+            generate_img_btn.click(
+                generate_img2img,
+                inputs=[prompt_img2img, input_img, strength, steps_img2img, guidance_img2img, width_img2img, height_img2img, seed_img2img, save_format_img2img],
+                outputs=output_img2img
+            )
 
-# Load model (use float32 for CPU)
-pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float32)
-pipe.to("cpu")
+    return demo
 
-# Image generation function
-def generate_image(prompt):
-    image = pipe(prompt).images[0]
-    return image
+# Launch Gradio WebUI
+web_ui = create_ui()
+url = web_ui.launch(share=True)
 
-# Gradio Interface
-demo = gr.Interface(fn=generate_image, inputs="text", outputs="image")
-demo.launch()
+# Automatically open the WebUI in a new browser tab
+webbrowser.open(url)