Spaces:
Sleeping
Sleeping
| import spaces | |
| import gradio as gr | |
| import torch | |
| from transformers import ( | |
| BlipProcessor, | |
| BlipForConditionalGeneration, | |
| AutoTokenizer, | |
| AutoModelForCausalLM, | |
| ) | |
| APP_TITLE = "StoryLens AI" | |
| STORY_STYLES = [ | |
| "Fantasy", | |
| "Adventure", | |
| "Mystery", | |
| "Sci-Fi", | |
| "Inspirational", | |
| "Comedy", | |
| "Historical", | |
| ] | |
| DEFAULT_STYLE = STORY_STYLES[0] | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| dtype = torch.float16 if device == "cuda" else torch.float32 | |
| blip_id = "Salesforce/blip-image-captioning-base" | |
| caption_processor = BlipProcessor.from_pretrained(blip_id) | |
| caption_model = BlipForConditionalGeneration.from_pretrained( | |
| blip_id, | |
| torch_dtype=dtype, | |
| ).to(device) | |
| llm_id = "Qwen/Qwen2.5-1.5B-Instruct" | |
| tokenizer = AutoTokenizer.from_pretrained(llm_id) | |
| story_model = AutoModelForCausalLM.from_pretrained( | |
| llm_id, | |
| torch_dtype=dtype, | |
| ).to(device) | |
| def generate_caption(image): | |
| inputs = caption_processor( | |
| image, | |
| return_tensors="pt" | |
| ).to(device, dtype) | |
| output = caption_model.generate( | |
| **inputs, | |
| max_new_tokens=30 | |
| ) | |
| return caption_processor.decode( | |
| output[0], | |
| skip_special_tokens=True | |
| ) | |
| def generate_story(caption, style, max_tokens): | |
| prompt = ( | |
| f"Create an engaging story inspired by this image description: '{caption}'. " | |
| f"Use the '{style}' style. " | |
| "The story should be creative, descriptive, and between 3 and 5 sentences." | |
| ) | |
| messages = [ | |
| { | |
| "role": "user", | |
| "content": prompt | |
| } | |
| ] | |
| inputs = tokenizer.apply_chat_template( | |
| messages, | |
| add_generation_prompt=True, | |
| return_tensors="pt", | |
| return_dict=False, | |
| ).to(device) | |
| outputs = story_model.generate( | |
| inputs, | |
| max_new_tokens=max_tokens, | |
| do_sample=True, | |
| temperature=0.8, | |
| top_p=0.9, | |
| pad_token_id=tokenizer.eos_token_id, | |
| ) | |
| return tokenizer.decode( | |
| outputs[0][inputs.shape[1]:], | |
| skip_special_tokens=True, | |
| ) | |
| def generate_output(image, style, length): | |
| if image is None: | |
| return "Please upload an image first.", "" | |
| image = image.convert("RGB") | |
| caption = generate_caption(image) | |
| story = generate_story( | |
| caption, | |
| style, | |
| length, | |
| ) | |
| return caption, story | |
| custom_theme = gr.themes.Default( | |
| primary_hue="blue", | |
| secondary_hue="blue", | |
| neutral_hue="slate", | |
| ) | |
| css = """ | |
| body{ | |
| background:#0f172a; | |
| } | |
| .gradio-container{ | |
| background:#0f172a !important; | |
| } | |
| footer{ | |
| display:none; | |
| } | |
| """ | |
| demo = gr.Interface( | |
| fn=generate_output, | |
| inputs=[ | |
| gr.Image( | |
| type="pil", | |
| label="Upload Image" | |
| ), | |
| gr.Dropdown( | |
| choices=STORY_STYLES, | |
| value=DEFAULT_STYLE, | |
| label="Story Style" | |
| ), | |
| gr.Slider( | |
| minimum=40, | |
| maximum=200, | |
| value=120, | |
| step=10, | |
| label="Story Length" | |
| ), | |
| ], | |
| outputs=[ | |
| gr.Textbox( | |
| label="Image Caption" | |
| ), | |
| gr.Textbox( | |
| label="Generated Story", | |
| lines=8 | |
| ), | |
| ], | |
| title=APP_TITLE, | |
| description=""" | |
| Upload any image and let AI first describe it, then transform that description | |
| into a creative story based on your selected writing style. | |
| """, | |
| theme=custom_theme, | |
| css=css, | |
| ) | |
| demo.launch(server_name="0.0.0.0", server_port=7860) |