Update app.py

app.py

@@ -1,88 +1,53 @@
 import torch
 import gradio as gr
-from transformers import pipeline, AutoModelForSeq2SeqLM, AutoTokenizer, AutoConfig
+from diffusers import StableDiffusionPipeline
 
-# List of summarization models
+# List of models for text-to-image generation
 model_names = [
-    "google/bigbird-pegasus-large-arxiv",
-    "facebook/bart-large-cnn",
-    "google/t5-v1_1-large",
-    "sshleifer/distilbart-cnn-12-6",
-    "allenai/led-base-16384",
-    "google/pegasus-xsum",
-    "togethercomputer/LLaMA-2-7B-32K"
+    "runwayml/stable-diffusion-v1-5",
+    "stabilityai/stable-diffusion-2-1",
+    "CompVis/stable-diffusion-v1-4"
 ]
 
-# Placeholder for the summarizer pipeline, tokenizer, and maximum tokens
-summarizer = None
-tokenizer = None
-max_tokens = None
+# Placeholder for pipeline
+pipe = None
 
-# Example text for summarization
-example_text = (
-    "Artificial intelligence (AI) is intelligence—perceiving, synthesizing, and inferring information—"
-    "demonstrated by machines, as opposed to intelligence displayed by non-human animals and humans. "
-    "Example tasks in which AI is employed include speech recognition, computer vision, language translation, "
-    "autonomous vehicles, and game playing. AI research has been defined as the field of study of intelligent "
-    "agents, which refers to any system that perceives its environment and takes actions that maximize its "
-    "chance of achieving its goals."
-)
-
-# Function to load the selected model
+# Function to load model
 def load_model(model_name):
-    global summarizer, tokenizer, max_tokens
+    global pipe
     try:
-        # Load the summarization pipeline with the selected model
-        summarizer = pipeline("summarization", model=model_name, torch_dtype=torch.float32)
-        tokenizer = AutoTokenizer.from_pretrained(model_name)
-        config = AutoConfig.from_pretrained(model_name)
-
-        # Set a reasonable default for max_tokens if not available
-        max_tokens = getattr(config, 'max_position_embeddings', 1024)
-
-        return f"Model {model_name} loaded successfully! Max tokens: {max_tokens}"
+        pipe = StableDiffusionPipeline.from_pretrained(model_name, torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32)
+        pipe = pipe.to("cuda" if torch.cuda.is_available() else "cpu")
+        return f"Model {model_name} loaded successfully!"
     except Exception as e:
-        return f"Failed to load model {model_name}. Error: {str(e)}"
-
-# Function to summarize the input text
-def summarize_text(input, min_length, max_length):
-    if summarizer is None:
-        return "No model loaded!"
+        return f"Failed to load model: {str(e)}"
 
+# Function to generate image from text
+def generate_image(prompt, guidance_scale):
+    if pipe is None:
+        return None, "Model not loaded yet!"
     try:
-        # Tokenize the input text and check the number of tokens
-        input_tokens = tokenizer.encode(input, return_tensors="pt")
-        num_tokens = input_tokens.shape[1]
-        if num_tokens > max_tokens:
-            return f"Error: Input exceeds the max token limit of {max_tokens}."
-
-        # Ensure min/max lengths are within bounds
-        min_summary_length = max(10, int(num_tokens * (min_length / 100)))
-        max_summary_length = min(max_tokens, int(num_tokens * (max_length / 100)))
-
-        # Summarize the input text
-        output = summarizer(input, min_length=min_summary_length, max_length=max_summary_length, truncation=True)
-        return output[0]['summary_text']
+        image = pipe(prompt, guidance_scale=guidance_scale).images[0]
+        return image, "Image generated successfully!"
     except Exception as e:
-        return f"Summarization failed: {str(e)}"
+        return None, f"Image generation failed: {str(e)}"
 
-# Gradio Interface
+# Gradio UI
 with gr.Blocks() as demo:
     with gr.Row():
-        model_dropdown = gr.Dropdown(choices=model_names, label="Choose a model", value="sshleifer/distilbart-cnn-12-6")
+        model_dropdown = gr.Dropdown(choices=model_names, label="Choose a model", value="runwayml/stable-diffusion-v1-5")
         load_button = gr.Button("Load Model")
+    load_status = gr.Textbox(label="Status", interactive=False)
 
-    load_message = gr.Textbox(label="Load Status", interactive=False)
-
-    min_length_slider = gr.Slider(minimum=0, maximum=100, step=1, label="Minimum Summary Length (%)", value=10)
-    max_length_slider = gr.Slider(minimum=0, maximum=100, step=1, label="Maximum Summary Length (%)", value=20)
-
-    input_text = gr.Textbox(label="Input text to summarize", lines=6, value=example_text)
-    summarize_button = gr.Button("Summarize Text")
-    output_text = gr.Textbox(label="Summarized text", lines=4)
-
-    load_button.click(fn=load_model, inputs=model_dropdown, outputs=load_message)
-    summarize_button.click(fn=summarize_text, inputs=[input_text, min_length_slider, max_length_slider],
-                           outputs=output_text)
+    with gr.Row():
+        prompt_input = gr.Textbox(label="Enter text prompt", placeholder="A futuristic city in the sky", lines=2)
+        guidance_slider = gr.Slider(minimum=1, maximum=20, value=7.5, step=0.5, label="Guidance Scale")
+    
+    generate_button = gr.Button("Generate Image")
+    output_image = gr.Image(label="Generated Image")
+    message = gr.Textbox(label="Generation Status", interactive=False)
+
+    load_button.click(fn=load_model, inputs=model_dropdown, outputs=load_status)
+    generate_button.click(fn=generate_image, inputs=[prompt_input, guidance_slider], outputs=[output_image, message])
 
 demo.launch()