sd app

app.py

@@ -1,51 +1,138 @@
 import gradio as gr
-from transformers import pipeline
+import torch
+from diffusers import StableDiffusionPipeline, StableDiffusionImg2ImgPipeline
+from PIL import Image
+import io
+import base64
 
-sentiment_analyzer = pipeline("sentiment-analysis")
-text_generator = pipeline("text-generation", model="gpt2")
-summarizer = pipeline("summarization")
+# Check if GPU is available
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
-def analyze_sentiment(text):
-    result = sentiment_analyzer(text)[0]
-    return f"{result['label']} ({result['score']:.2%})"
+# Load text-to-image pipeline
+pipe_text_to_image = StableDiffusionPipeline.from_pretrained(
+    "runwayml/stable-diffusion-v1-5",
+    torch_dtype=torch.float16 if device == "cuda" else torch.float32
+).to(device)
 
-def generate_text(prompt):
-    result = text_generator(prompt, max_length=100, num_return_sequences=1)
-    return result[0]['generated_text']
+# Load image-to-image pipeline
+pipe_image_to_image = StableDiffusionImg2ImgPipeline.from_pretrained(
+    "runwayml/stable-diffusion-v1-5",
+    torch_dtype=torch.float16 if device == "cuda" else torch.float32
+).to(device)
 
-def summarize_text(text):
-    result = summarizer(text, max_length=150, min_length=30, do_sample=False)
-    return result[0]['summary_text']
+def generate_image(prompt, negative_prompt="", steps=20, guidance=7.5):
+    """Generate image from text prompt"""
+    if not prompt.strip():
+        return None, "Please enter a prompt"
+    
+    try:
+        with torch.autocast(device):
+            image = pipe_text_to_image(
+                prompt=prompt,
+                negative_prompt=negative_prompt,
+                num_inference_steps=steps,
+                guidance_scale=guidance
+            ).images[0]
+        return image, "Image generated successfully!"
+    except Exception as e:
+        return None, f"Error: {str(e)}"
 
-with gr.Blocks(title="AI Toolkit") as demo:
-    gr.Markdown("# 🚀 AI Toolkit")
+def modify_image(image, prompt, strength=0.75, steps=20, guidance=7.5):
+    """Modify existing image with prompt"""
+    if image is None:
+        return None, "Please upload an image first"
+    if not prompt.strip():
+        return None, "Please enter a modification prompt"
     
-    with gr.Tab("Sentiment Analysis"):
-        with gr.Row():
-            with gr.Column():
-                sentiment_input = gr.Textbox(label="Enter text for sentiment analysis", lines=3)
-                sentiment_button = gr.Button("Analyze Sentiment")
-            with gr.Column():
-                sentiment_output = gr.Textbox(label="Result")
+    try:
+        # Resize image to appropriate dimensions
+        width, height = image.size
+        max_size = 512
+        if width > max_size or height > max_size:
+            image.thumbnail((max_size, max_size), Image.Resampling.LANCZOS)
+        
+        with torch.autocast(device):
+            result = pipe_image_to_image(
+                prompt=prompt,
+                image=image,
+                strength=strength,
+                num_inference_steps=steps,
+                guidance_scale=guidance
+            ).images[0]
+        return result, "Image modified successfully!"
+    except Exception as e:
+        return None, f"Error: {str(e)}"
+
+with gr.Blocks(title="AI Image Modifier") as demo:
+    gr.Markdown("# 🎨 AI Image Generator & Modifier")
+    gr.Markdown("Create new images or modify existing ones using text prompts!")
     
-    with gr.Tab("Text Generation"):
+    with gr.Tab("Generate New Image"):
         with gr.Row():
             with gr.Column():
-                generation_input = gr.Textbox(label="Enter prompt", lines=2)
-                generation_button = gr.Button("Generate Text")
+                gen_prompt = gr.Textbox(
+                    label="Describe your image",
+                    placeholder="A beautiful sunset over mountains, digital art...",
+                    lines=2
+                )
+                gen_negative = gr.Textbox(
+                    label="What to avoid (optional)",
+                    placeholder="blurry, low quality, distorted...",
+                    lines=1
+                )
+                with gr.Row():
+                    gen_steps = gr.Slider(1, 50, value=20, label="Steps")
+                    gen_guidance = gr.Slider(1, 20, value=7.5, label="Guidance Scale")
+                
+                gen_button = gr.Button("Generate Image", variant="primary")
+            
             with gr.Column():
-                generation_output = gr.Textbox(label="Generated Text", lines=5)
+                gen_output = gr.Image(label="Generated Image")
+                gen_status = gr.Textbox(label="Status", interactive=False)
     
-    with gr.Tab("Text Summarization"):
+    with gr.Tab("Modify Existing Image"):
         with gr.Row():
             with gr.Column():
-                summary_input = gr.Textbox(label="Enter text to summarize", lines=5)
-                summary_button = gr.Button("Summarize")
+                mod_input_image = gr.Image(label="Upload Image", type="pil")
+                mod_prompt = gr.Textbox(
+                    label="How do you want to modify it?",
+                    placeholder="Make it look like winter, change style to oil painting...",
+                    lines=2
+                )
+                with gr.Row():
+                    mod_strength = gr.Slider(0.1, 1.0, value=0.75, label="Modification Strength")
+                    mod_steps = gr.Slider(1, 50, value=20, label="Steps")
+                    mod_guidance = gr.Slider(1, 20, value=7.5, label="Guidance Scale")
+                
+                mod_button = gr.Button("Modify Image", variant="primary")
+            
             with gr.Column():
-                summary_output = gr.Textbox(label="Summary", lines=3)
+                mod_output = gr.Image(label="Modified Image")
+                mod_status = gr.Textbox(label="Status", interactive=False)
+    
+    # Examples
+    gr.Markdown("### Example Prompts:")
+    gr.Examples(
+        examples=[
+            ["A majestic dragon flying over a medieval castle, fantasy art, highly detailed"],
+            ["A cyberpunk cityscape at night, neon lights, raining, futuristic"],
+            ["A cute corgi puppy wearing a superhero cape, cartoon style"],
+            ["An astronaut riding a horse on Mars, photorealistic"]
+        ],
+        inputs=gen_prompt
+    )
+    
+    # Connect functions
+    gen_button.click(
+        generate_image,
+        inputs=[gen_prompt, gen_negative, gen_steps, gen_guidance],
+        outputs=[gen_output, gen_status]
+    )
     
-    sentiment_button.click(analyze_sentiment, inputs=sentiment_input, outputs=sentiment_output)
-    generation_button.click(generate_text, inputs=generation_input, outputs=generation_output)
-    summary_button.click(summarize_text, inputs=summary_input, outputs=summary_output)
+    mod_button.click(
+        modify_image,
+        inputs=[mod_input_image, mod_prompt, mod_strength, mod_steps, mod_guidance],
+        outputs=[mod_output, mod_status]
+    )
 
-demo.launch()
+demo.launch(share=True)