Update app.py

app.py

@@ -1,37 +1,30 @@
 import gradio as gr
-from transformers import GPT2LMHeadModel, GPT2Tokenizer
+from diffusers import StableDiffusionPipeline
+import torch
 
-# Load pre-trained GPT-2 model and tokenizer
-model_name = "gpt2-medium"  # You can use "gpt2", "gpt2-medium", "gpt2-large", or "gpt2-xl"
-model = GPT2LMHeadModel.from_pretrained(model_name)
-tokenizer = GPT2Tokenizer.from_pretrained(model_name)
+# Load a lightweight stable diffusion model
+model_id = "CompVis/stable-diffusion-v1-4"
+pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
+pipe = pipe.to("cuda")  # Use GPU for faster generation, or .to("cpu") if no GPU is available
 
-# Function to generate text based on input
-def generate_story(prompt, choice1="", choice2="", choice3=""):
-    # Combine the prompt with the chosen path
-    full_prompt = prompt + "\n" + choice1 + "\n" + choice2 + "\n" + choice3
-    inputs = tokenizer.encode(full_prompt, return_tensors="pt")
-    outputs = model.generate(inputs, max_length=150, num_return_sequences=1, pad_token_id=tokenizer.eos_token_id)
-    text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return text
+def generate_character(name, gender, trait, superpower):
+    description = f"A {gender} character named {name}, who is {trait}. Their superpower is {superpower}."
+    # Generate an image based on the description
+    image = pipe(description).images[0]
+    return image
 
 # Gradio interface
-def adventure_game(prompt, choice1, choice2, choice3):
-    story = generate_story(prompt, choice1, choice2, choice3)
-    return story
-
-# Create Gradio interface
 iface = gr.Interface(
-    fn=adventure_game, 
+    fn=generate_character,
     inputs=[
-        gr.Textbox(label="Adventure Start", placeholder="Enter the beginning of your adventure..."),
-        gr.Textbox(label="Choice 1", placeholder="Enter your first choice..."),
-        gr.Textbox(label="Choice 2", placeholder="Enter your second choice..."),
-        gr.Textbox(label="Choice 3", placeholder="Enter your third choice...")
-    ], 
-    outputs="text",
-    title="Text-Based Adventure Game Generator",
-    description="Create your own adventure by entering the beginning of a story and making choices to continue the adventure."
+        gr.Textbox(label="Name", placeholder="Enter the character's name"),
+        gr.Radio(label="Gender", choices=["male", "female", "non-binary"]),
+        gr.Textbox(label="Personality Trait", placeholder="e.g., brave, clever, kind"),
+        gr.Textbox(label="Superpower", placeholder="e.g., flying, invisibility, super strength")
+    ],
+    outputs="image",
+    title="Lightweight Character Generator",
+    description="Enter the details of your character and generate an image!"
 )
 
 iface.launch()